import gradio as gr
from ultralytics import YOLO
import cv2
import tempfile

# Carregar o modelo YOLOv8
model = YOLO("yolov8n.pt")  # Você pode usar 'yolov8s.pt', 'yolov8m.pt', etc., dependendo do seu requisito

def detect_objects(video_path):
    # Abrir o vídeo de entrada
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise ValueError("Não foi possível abrir o vídeo.")

    # Criar um arquivo temporário para salvar o vídeo processado
    output_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, 30.0, (int(cap.get(3)), int(cap.get(4))))

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break

        # Realizar a detecção de objetos
        results = model(frame)
        annotated_frame = results[0].plot()  # Obter o frame com as detecções

        # Escrever o frame processado no vídeo de saída
        out.write(annotated_frame)

    # Liberar os recursos
    cap.release()
    out.release()

    # Retornar o caminho do vídeo de entrada e do vídeo processado
    return video_path, output_path

# Interface Gradio
iface = gr.Interface(
    fn=detect_objects,
    inputs=gr.Video(label="Vídeo de Entrada"),
    outputs=[
        gr.Video(label="Vídeo de Entrada"),
        gr.Video(label="Vídeo com Detecção de Objetos")
    ],
    title="Detecção de Objetos com YOLOv8",
    description="Carregue um vídeo curto para detectar objetos usando YOLOv8. Os vídeos de entrada e saída serão exibidos lado a lado."
)

# Layout personalizado para exibir os vídeos lado a lado
iface = gr.Interface(
    fn=detect_objects,
    inputs=gr.Video(label="Vídeo de Entrada"),
    outputs=gr.Row(
        gr.Video(label="Vídeo de Entrada"),
        gr.Video(label="Vídeo com Detecção de Objetos")
    ),
    title="Detecção de Objetos com YOLOv8",
    description="Carregue um vídeo curto para detectar objetos usando YOLOv8. Os vídeos de entrada e saída serão exibidos lado a lado."
)

# Iniciar a interface
iface.launch()