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

# Initialize models
model_path = hf_hub_download(repo_id="brainwavecollective/yolov8n-rubber-duck-detector", filename="yolov8n_rubberducks.pt")
duck_model = YOLO(model_path)
standard_model = YOLO('yolov8n.pt')

def process_image(image, model, is_standard_model=True):
    results = model(image)
    processed_image = image.copy()
    
    for r in results:
        boxes = r.boxes
        for box in boxes:
            cls = int(box.cls[0])
            class_name = model.names[cls]
            
            # For standard model, only show teddy bears
            if is_standard_model and class_name != "teddy bear":
                continue
                
            x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy())
            conf = float(box.conf[0])
            
            # Rename class to "rubber duck"
            display_name = "rubber duck" if not is_standard_model else class_name
            
            cv2.rectangle(processed_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
            label = f"{display_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):
    image = np.array(input_image)
    
    standard_image = process_image(image, standard_model, is_standard_model=True)
    duck_image = process_image(image, duck_model, is_standard_model=False)
    
    height, width = image.shape[:2]
    gap = 20  # Add a 20-pixel gap between images
    
    # Create canvas with extra width for the gap
    canvas = np.zeros((height, width * 2 + gap, 3), dtype=np.uint8)
    
    # Place images with gap in between
    canvas[:, :width] = standard_image
    canvas[:, width + gap:] = duck_image
    
    # Fill gap with white or gray color (optional)
    canvas[:, width:width + gap] = [128, 128, 128]  # Gray gap
    
    # Add model labels
    cv2.putText(canvas, "YOLOv8", (10, 30), 
                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
    cv2.putText(canvas, "YOLOv8 Fine Tune", (width + gap + 10, 30), 
                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
    
    return canvas

# Create Gradio interface with stacked layout
with gr.Blocks() as iface:
    gr.Markdown("# Improved YOLO Rubber Duck Detection")
    gr.Markdown("Compare standard YOLOv8 (left) with [Fine-tuned YOLOv8 Rubber Duck Detector](https://huggingface.co/brainwavecollective/yolov8n-rubber-duck-detector) (right)")
    
    with gr.Column():
        # First define the input and output components
        input_image = gr.Image(type="pil", label="Input Image", height=200)
        submit_btn = gr.Button("Compare Models")
        output_image = gr.Image(type="numpy", label="Comparison Result", height=400)
        
        # Then add Examples section that references them
        gr.Examples(
            examples=[["test_image.jpg"]],
            inputs=input_image,
            outputs=output_image,
            fn=compare_models,
            cache_examples=True
        )
    
    submit_btn.click(
        fn=compare_models,
        inputs=input_image,
        outputs=output_image,
    )

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