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

# Cargar el modelo YOLO
model = YOLO('best.pt')  # Asegúrate de que 'best.pt' es el camino correcto a tu modelo

def plot_bboxes(image, results, labels=None, colors=None, conf=0.2, score=True):
    """
    Dibuja las bounding boxes en la imagen con base en los resultados de YOLO.
    """
    image_copy = image.copy()
    if len(results) > 0 and hasattr(results[0], 'boxes'):
        for box in results[0].boxes:
            x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)
            cls = int(box.cls[0].cpu().numpy())
            conf_val = float(box.conf[0].cpu().numpy())

            if conf_val < conf:
                continue

            color = colors[cls % len(colors)] if colors else (0, 255, 0)
            label = labels.get(cls, f"Clase {cls}") if labels else f"Clase {cls}"
            if score:
                label += f" ({conf_val:.2f})"
            cv2.rectangle(image_copy, (x1, y1), (x2, y2), color, 2)
            cv2.putText(image_copy, label, (x1, y1 - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
    return image_copy

def process_video(video):
    """
    Procesa un video en tiempo real desde la webcam, aplicando YOLO a cada frame.
    """
    cap = cv2.VideoCapture(video)  # Recibe la entrada de video (webcam)
    ret, frame = cap.read()
    if not ret:
        return None
    frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)  # Convertir de RGB a BGR
    results = model.predict(frame_bgr, conf=0.5)  # Realizar la predicción
    annotated_frame = plot_bboxes(
        frame_bgr,
        results,
        labels=model.names,
        conf=0.3,
        score=True
    )
    annotated_frame_rgb = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)  # Convertir de nuevo a RGB
    cap.release()
    return annotated_frame_rgb

# Crear la interfaz Gradio usando gr.Video
iface = gr.Interface(
    fn=process_video,
    inputs=gr.Video(),  # Entrada de video
    outputs=gr.Image(type="numpy"),  # Salida de la imagen procesada
    title="Trashformers: Clasificación en Tiempo Real",
    description="Detecta y clasifica residuos con YOLO directamente desde tu webcam."
)

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