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

# Safe load method to handle custom YOLO class during deserialization
def safe_load_yolo_model(path):
    # Add necessary safe globals to allow the detection model class during loading
    torch.serialization.add_safe_globals([torch, 'ultralytics.nn.tasks.DetectionModel'])
    return YOLO(path)

# Load YOLO models
model_yolo11 = safe_load_yolo_model('./data/yolo11n.pt')
model_best = safe_load_yolo_model('./data/best.pt')

def process_video(video):
    # Read video input
    cap = cv2.VideoCapture(video.name)
    fps = cap.get(cv2.CAP_PROP_FPS)
    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    # Create a VideoWriter object to save the output video
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # Codec for .mp4
    out = cv2.VideoWriter('output_video.mp4', fourcc, fps, (frame_width, frame_height))

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        
        # Use both YOLO models for detection
        results_yolo11 = model_yolo11(frame)
        results_best = model_best(frame)
        
        # Combine the results from both models
        # For simplicity, we will overlay bounding boxes and labels from both models
        for result in results_yolo11:
            boxes = result.boxes
            for box in boxes:
                x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                label = f"YOLOv11: {box.cls[0]} - {box.conf[0]:.2f}"
                cv2.putText(frame, label, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)

        for result in results_best:
            boxes = result.boxes
            for box in boxes:
                x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
                cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 0), 2)
                label = f"Best: {box.cls[0]} - {box.conf[0]:.2f}"
                cv2.putText(frame, label, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)

        # Write the processed frame to the output video
        out.write(frame)

    cap.release()
    out.release()

    return 'output_video.mp4'

# Gradio interface
iface = gr.Interface(fn=process_video, inputs=gr.Video(), outputs=gr.Video(), live=True)

# Launch the app
iface.launch()