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

# Carregar o modelo YOLOv8
model = YOLO("yolov8n.pt")  # Use o modelo mais leve (nano)

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.")
    
    # Configurar a saída de vídeo
    output_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) // 2)  # Reduzir a largura pela metade
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) // 2)  # Reduzir a altura pela metade
    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
    
    frame_skip = 2  # Pular 1 frame a cada 2 (processar apenas metade dos frames)
    frame_count = 0
    max_frames = int(fps * 5)  # Processar no máximo 5 segundos de vídeo
    
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret or frame_count >= max_frames:
            break
        
        # Pular frames intermediários
        if frame_count % frame_skip != 0:
            frame_count += 1
            continue
        
        # Redimensionar o frame para metade da resolução original
        resized_frame = cv2.resize(frame, (width, height))
        
        # Realizar a detecção de objetos
        results = model(resized_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)
        
        frame_count += 1
    
    # 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

# Usar gr.Blocks para criar uma interface personalizada
with gr.Blocks() as demo:
    gr.Markdown("# Detecção de Objetos com YOLOv8")
    gr.Markdown("Carregue um vídeo curto para detectar objetos usando YOLOv8. Os vídeos de entrada e saída serão exibidos lado a lado.")
    
    with gr.Row():
        input_video = gr.Video(label="Vídeo de Entrada")
        output_video = gr.Video(label="Vídeo com Detecção de Objetos")
    
    detect_button = gr.Button("Detectar Objetos")
    
    def process_and_display(video_path):
        original_video, processed_video = detect_objects(video_path)
        return original_video, processed_video
    
    detect_button.click(
        fn=process_and_display,
        inputs=input_video,
        outputs=[input_video, output_video]
    )

# Iniciar a interface
demo.launch()