import streamlit as st import os import cv2 import torch from tqdm import tqdm from ultralytics import YOLO import pandas as pd import numpy as np from PIL import Image from VideoProcessor import MediaProcessor, process_media # Создание папок для загрузки и обработки файлов def create_folders(upload_folder="uploaded_files", processed_folder="processed_files"): if not os.path.exists(upload_folder): os.makedirs(upload_folder) if not os.path.exists(processed_folder): os.makedirs(processed_folder) # Функция для загрузки файлов def save_uploaded_file(uploaded_file, folder_name="uploaded_files"): file_path = os.path.join(folder_name, uploaded_file.name) with open(file_path, "wb") as f: f.write(uploaded_file.getbuffer()) return file_path # Функция для получения списка всех файлов в папке def get_all_files(folder_name="processed_files"): return os.listdir(folder_name) # Функция для отображения файлов с центровкой def display_file(selected_file, folder_name="processed_files"): file_path = os.path.join(folder_name, selected_file) if selected_file.endswith('.mp4'): st.video(file_path) else: st.image(file_path, use_column_width=True) def exclude_processed_files(file_list, processed_files): return [file for file in file_list if os.path.basename(file.name) not in processed_files] # Основная функция приложения def main(processor): variants = [] processed_files = [] # Массив для хранения уже обработанных файлов # Создание папок для загрузки и обработки файлов create_folders() # Заголовок приложения st.title("Загрузите фото и видео, затем выберите файл из списка") # Загрузка файлов uploaded_files = st.file_uploader("Загрузите фото и видео", accept_multiple_files=True) if uploaded_files: input_paths = [] # Исключение уже обработанных файлов new_files = exclude_processed_files(uploaded_files, processed_files) for uploaded_file in new_files: file_path = save_uploaded_file(uploaded_file) input_paths.append(file_path) if input_paths: st.toast(f"Файлы загружены", icon="🟢") imgs, vids = process_media(input_paths, processor) # Получение реальных названий файлов с расширениями, но без папки variants.extend([os.path.basename(i) for i in imgs]) variants.extend([os.path.basename(i) for i in vids]) st.toast(f"Файлы обработаны", icon="🟢") # Добавление обработанных файлов в processed_files processed_files.extend([os.path.basename(i) for i in imgs]) processed_files.extend([os.path.basename(i) for i in vids]) # Поле для выбора файла из выпадающего списка selected_file = st.selectbox("Выберите файл", variants) # Центровка и отображение выбранного файла if selected_file: st.markdown( """ """, unsafe_allow_html=True ) st.markdown('
', unsafe_allow_html=True) display_file(selected_file) st.markdown('
', unsafe_allow_html=True) # Запуск приложения if __name__ == "__main__": model_path = 'trained_y8m.pt' # Укажите путь к модели processor = MediaProcessor('processed_files', model_path, batch_size=16) main(processor)