import gradio as gr
from ultralytics import YOLO
from PIL import Image

# Load BOTH models
base_model = YOLO("yolo11s.pt")
finetuned_model = YOLO("WildLife3-best.onnx")

def compare_models(image, conf_threshold):
    """Compare base model vs fine-tuned model"""
    
    if image is None:
        empty_img = Image.new('RGB', (640, 480), color='white')
        return empty_img, "Please upload an image", empty_img, "Please upload an image"
    
    # Base model prediction
    base_results = base_model.predict(source=image, conf=conf_threshold)
    base_annotated = Image.fromarray(base_results[0].plot()[..., ::-1])
    
    base_boxes = base_results[0].boxes
    base_detections = []
    for box in base_boxes:
        cls = int(box.cls[0])
        conf = float(box.conf[0])
        name = base_results[0].names[cls]
        base_detections.append(f"{name}: {conf:.2%}")
    base_text = "\n".join(base_detections) if base_detections else "No detections"
    
    # Fine-tuned model prediction
    ft_results = finetuned_model.predict(source=image, conf=conf_threshold)
    ft_annotated = Image.fromarray(ft_results[0].plot()[..., ::-1])
    
    ft_boxes = ft_results[0].boxes
    ft_detections = []
    for box in ft_boxes:
        cls = int(box.cls[0])
        conf = float(box.conf[0])
        name = ft_results[0].names[cls]
        ft_detections.append(f"{name}: {conf:.2%}")
    ft_text = "\n".join(ft_detections) if ft_detections else "No detections"
    
    return base_annotated, base_text, ft_annotated, ft_text

with gr.Blocks(title="Wildlife Detector Comparison") as demo:
    gr.Markdown(
        """
        # 🦁 Model Comparison: Base YOLO11s vs Fine-Tuned Wildlife Detector
        
        **Left:** Pre-trained YOLO11s (80 COCO classes - general objects)  
        **Right:** Fine-tuned YOLO11s (Wildlife species detector)
        
        See how fine-tuning improves wildlife detection!
        
        **The model was trained on these 20 wildlife species:**
        
        Snow Leopard, Tiger, Leopard, Gorilla, Zebra, Peacock, Panda, Hyena, Pig, Horse, Dog, Donkey, Elephant, Fox, Hippopotamus, Kangaroo, Lion, Sheep, Wolf
        """
    )
    
    with gr.Row():
        image_input = gr.Image(type="pil", label="Upload Animal Image")
        conf_slider = gr.Slider(0.1, 1.0, value=0.5, label="Confidence Threshold")
    
    detect_btn = gr.Button("Compare Models", variant="primary")
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("### 📦 Base YOLO11s (COCO)")
            base_output = gr.Image(label="Base Model Detection")
            base_text = gr.Textbox(label="Base Model Detections", lines=8)
        
        with gr.Column():
            gr.Markdown("### 🎯 Fine-Tuned (Wildlife)")
            ft_output = gr.Image(label="Fine-Tuned Detection")
            ft_text = gr.Textbox(label="Fine-Tuned Detections", lines=8)
    
    detect_btn.click(
        fn=compare_models,
        inputs=[image_input, conf_slider],
        outputs=[base_output, base_text, ft_output, ft_text]
    )

demo.launch()