app.py

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(
            """
            <style>
            .centered {
                display: flex;
                justify-content: center;
            }
            </style>
            """,
            unsafe_allow_html=True
        )
        st.markdown('<div class="centered">', unsafe_allow_html=True)
        display_file(selected_file)
        st.markdown('</div>', unsafe_allow_html=True)

# Запуск приложения
if __name__ == "__main__":
    model_path = 'trained_y8m.pt'  # Укажите путь к модели
    processor = MediaProcessor('processed_files', model_path, batch_size=16)

    main(processor)