import gradio as gr
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
from ultralytics import YOLO
import time

# Initialize YOLO model
def load_model(model_size="medium"):
    model_map = {
        "nano": "yolov8n.pt",
        "small": "yolov8s.pt",
        "medium": "yolov8m.pt",
        "large": "yolov8l.pt",
        "xlarge": "yolov8x.pt"
    }
    try:
        model = YOLO(model_map[model_size])
        return model, f"✅ {model_size.capitalize()} model loaded successfully"
    except Exception as e:
        return None, f"❌ Error loading model: {str(e)}"

# Enhanced detection function
def detect_objects(image, model_size="medium", conf_threshold=0.35, iou_threshold=0.45):
    # Load model (load fresh each time to change model size dynamically)
    model, status = load_model(model_size)
    if model is None:
        return image, status
    
    # Convert image to numpy array if needed
    if isinstance(image, np.ndarray):
        img = image
    else:
        img = np.array(image)
    
    # Run inference with improved parameters
    results = model.predict(
        source=img,
        conf=conf_threshold,
        iou=iou_threshold,
        imgsz=640,  # Increased from default 320 for better detection
        augment=True,  # Enable test-time augmentation
        visualize=False  # Disable visualization for speed
    )
    
    # Extract detailed detection information
    detections = []
    for r in results:
        boxes = r.boxes.cpu().numpy()
        for i, box in enumerate(boxes):
            detections.append({
                "class": model.names[int(box.cls[0])],
                "confidence": float(box.conf[0]),
                "bbox": box.xyxy[0].tolist(),
                "id": i+1
            })
    
    # Sort detections by confidence (highest first)
    detections.sort(key=lambda x: x["confidence"], reverse=True)
    
    # Create annotated image with enhanced visualization
    annotated_img = results[0].plot(
        line_width=2,
        font_size=14,
        pil=True
    )
    
    # Generate detailed detection text
    detection_text = "Detection Results:\n"
    if not detections:
        detection_text = "No objects detected"
    else:
        for det in detections[:20]:  # Show top 20 detections
            detection_text += (
                f"{det['id']}. {det['class'].upper()} "
                f"(Confidence: {det['confidence']:.1%})\n"
                f"   Bounding Box: {[int(x) for x in det['bbox']]}\n"
            )
    
    return annotated_img, detection_text

# Gradio interface
with gr.Blocks(title="Enhanced Object Detector") as demo:
    gr.Markdown("# 🚀 Enhanced Object Detection with YOLOv8")
    gr.Markdown("### Get comprehensive detections with adjustable parameters")
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(label="Input Image", type="pil")
            with gr.Accordion("Detection Settings", open=False):
                model_size = gr.Dropdown(
                    ["nano", "small", "medium", "large", "xlarge"],
                    value="medium",
                    label="Model Size"
                )
                conf_slider = gr.Slider(
                    minimum=0.1, maximum=0.9, step=0.05,
                    value=0.35, label="Confidence Threshold"
                )
                iou_slider = gr.Slider(
                    minimum=0.1, maximum=0.9, step=0.05,
                    value=0.45, label="IOU Threshold"
                )
            detect_btn = gr.Button("🔍 Detect Objects", variant="primary")

        with gr.Column():
            image_output = gr.Image(label="Detected Objects", interactive=False)
            text_output = gr.Textbox(
                label="Detection Details",
                lines=15,
                max_lines=20,
                interactive=False
            )
    
    detect_btn.click(
        fn=detect_objects,
        inputs=[image_input, model_size, conf_slider, iou_slider],
        outputs=[image_output, text_output]
    )

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