import gradio as gr
import torch
from ultralytics import YOLO
import numpy as np
from PIL import Image
import os

# Load YOLOv8 Small model
print("Loading YOLOv8 Small model...")
model = YOLO('yolov8s.pt')
print("Model loaded successfully!")

def predict(image, confidence_threshold, iou_threshold):
    """
    Run YOLOv8 inference on the input image
    """
    try:
        # Run inference
        results = model(
            image, 
            conf=confidence_threshold,
            iou=iou_threshold
        )
        
        # Get the first result (single image)
        result = results[0]
        
        # Create annotated image
        annotated_image = result.plot()
        
        # Convert BGR to RGB for proper display
        annotated_image = Image.fromarray(annotated_image[..., ::-1])
        
        # Get detection info
        detections = []
        if result.boxes is not None:
            for box in result.boxes:
                class_id = int(box.cls[0])
                class_name = model.names[class_id]
                confidence = float(box.conf[0])
                detections.append(f"{class_name}: {confidence:.2f}")
        
        detection_text = "\n".join(detections) if detections else "No objects detected"
        
        return annotated_image, detection_text
        
    except Exception as e:
        return None, f"Error: {str(e)}"

# Create Gradio interface
with gr.Blocks(title="YOLOv8 Object Detection") as demo:
    gr.Markdown("# 🔍 YOLOv8 Small Object Detection")
    gr.Markdown("Upload an image to detect objects using YOLOv8s model")
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="Upload Image")
            confidence_slider = gr.Slider(
                minimum=0.1, 
                maximum=1.0, 
                value=0.5, 
                step=0.05,
                label="Confidence Threshold"
            )
            iou_slider = gr.Slider(
                minimum=0.1, 
                maximum=1.0, 
                value=0.5, 
                step=0.05,
                label="IoU Threshold"
            )
            submit_btn = gr.Button("Detect Objects", variant="primary")
            
        with gr.Column():
            output_image = gr.Image(type="pil", label="Detection Results")
            output_text = gr.Textbox(label="Detected Objects", lines=10)
    
    # Connect the function
    submit_btn.click(
        fn=predict,
        inputs=[input_image, confidence_slider, iou_slider],
        outputs=[output_image, output_text]
    )
    
    # Add examples
    gr.Examples(
        examples=[
            ["https://ultralytics.com/images/bus.jpg", 0.5, 0.5],
        ],
        inputs=[input_image, confidence_slider, iou_slider],
        outputs=[output_image, output_text],
        fn=predict,
        cache_examples=False
    )

if __name__ == "__main__":
    demo.launch()