import cv2
import gradio as gr
from transformers import pipeline

# Cargar el modelo de detección de objetos de Hugging Face (DETR)
detector = pipeline("object-detection", model="facebook/detr-resnet-50")

def process_video(video_path):
    """
    Procesa un video y devuelve el máximo número detectado de personas, bicicletas y motos en un fotograma.
    """
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return {"person": 0, "bicycle": 0, "motorcycle": 0}
    
    max_counts = {"person": 0, "bicycle": 0, "motorcycle": 0}
    
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        
        # Convertir el frame de BGR a RGB (requerido por el modelo)
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        
        # Realizar la detección de objetos
        results = detector(frame_rgb)
        
        # Contar objetos detectados en el frame actual (usamos un umbral de confianza)
        frame_counts = {"person": 0, "bicycle": 0, "motorcycle": 0}
        for detection in results:
            if detection["score"] < 0.7:
                continue
            label = detection["label"].lower()
            if label in frame_counts:
                frame_counts[label] += 1
        
        # Actualizar el conteo máximo si en este frame se detecta más
        for key in frame_counts:
            if frame_counts[key] > max_counts[key]:
                max_counts[key] = frame_counts[key]
    
    cap.release()
    return max_counts

# Crear la interfaz de Gradio para el Space
iface = gr.Interface(
    fn=process_video,
    inputs=gr.Video(label="Sube tu video"),
    outputs="json",
    title="Detección de Objetos en Video",
    description="Carga un video y detecta cuántas personas, bicicletas y motos aparecen usando modelos de Hugging Face."
)

if __name__ == "__main__":
    iface.launch()