app.py

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

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

def process_image(image, model):
    results = model(image)
    processed_image = image.copy()
    
    for r in results:
        boxes = r.boxes
        for box in boxes:
            # Get box coordinates and confidence
            x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy())
            conf = float(box.conf[0])
            cls = int(box.cls[0])
            class_name = model.names[cls]
            
            # Draw box and label
            cv2.rectangle(processed_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
            label = f"{class_name} ({conf:.2f})"
            cv2.putText(processed_image, label, (x1, y1-10), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
    
    return processed_image

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 = process_image(image, duck_model)
    standard_image = process_image(image, standard_model)
    
    # 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 model labels
    cv2.putText(canvas, "Duck Detector", (10, 30), 
                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)
    
    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 Duck Detector with standard YOLOv8 model",
    examples=[["test_image.jpg"]],
    cache_examples=True
)

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