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

# Initialize models
duck_model = YOLO('https://huggingface.co/brainwavecollective/yolo8n-rubber-duck-detector/resolve/main/yolov8n_rubberducks4.pt')
standard_model = YOLO('yolov8n.pt')

def calculate_iou(box1, box2):
    x1 = max(box1[0], box2[0])
    y1 = max(box1[1], box2[1])
    x2 = min(box1[2], box2[2])
    y2 = min(box1[3], box2[3])
    intersection = max(0, x2 - x1) * max(0, y2 - y1)
    
    box1_area = (box1[2] - box1[0]) * (box1[3] - box1[1])
    box2_area = (box2[2] - box2[0]) * (box2[3] - box2[1])
    union = box1_area + box2_area - intersection
    
    return intersection / union if union > 0 else 0

def process_image(image, model, is_duck_model=True):
    results = model(image, conf=0.4)
    
    valid_boxes = []
    for r in results:
        for box in r.boxes:
            class_name = model.names[int(box.cls[0])]
            # For both models, only show teddy bear class
            if class_name == "teddy bear":
                valid_boxes.append({
                    'coords': box.xyxy[0].cpu().numpy().tolist(),
                    'confidence': float(box.conf[0])
                })
    
    # Filter overlapping boxes
    filtered_boxes = []
    for i, box in enumerate(valid_boxes):
        should_add = True
        for existing_box in filtered_boxes:
            if calculate_iou(box['coords'], existing_box['coords']) > 0.5:
                if box['confidence'] <= existing_box['confidence']:
                    should_add = False
                    break
        if should_add:
            filtered_boxes.append(box)
    
    # Draw boxes and labels
    processed_image = image.copy()
    for box in filtered_boxes:
        x1, y1, x2, y2 = map(int, box['coords'])
        cv2.rectangle(processed_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
        label = f"{'Duck' if is_duck_model else 'Standard'} ({box['confidence']:.2f})"
        cv2.putText(processed_image, label, (x1, y1-10), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
    
    return processed_image, len(filtered_boxes)

def compare_models(input_image):
    # Convert from Gradio's PIL image to OpenCV format
    image = np.array(input_image)
    
    # Process with both models
    duck_image, duck_detections = process_image(image, duck_model, True)
    standard_image, std_detections = process_image(image, standard_model, False)
    
    # Create side-by-side comparison
    height, width = image.shape[:2]
    canvas = np.zeros((height, width * 2, 3), dtype=np.uint8)
    
    # Place images side by side
    canvas[:, :width] = duck_image
    canvas[:, width:] = standard_image
    
    # Add labels for each model
    cv2.putText(canvas, "Rubber Duck YOLO", (10, 30), 
                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
    cv2.putText(canvas, f"Detections: {duck_detections}", (10, 60), 
                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
    
    cv2.putText(canvas, "Standard YOLOv8", (width + 10, 30), 
                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
    cv2.putText(canvas, f"Detections: {std_detections}", (width + 10, 60), 
                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
    
    return canvas

# Create Gradio interface
iface = gr.Interface(
    fn=compare_models,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="numpy"),
    title="YOLO Model Comparison",
    description="Compare Rubber Duck YOLO detector with standard YOLOv8 model",
    examples=[["example1.jpg"], ["example2.jpg"]],  # Add your example images here
    cache_examples=True
)

# Launch the interface
if __name__ == "__main__":
    iface.launch(share=True)