import torch
import cv2
import numpy as np
import gradio as gr
from PIL import Image

# Device configuration
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Load YOLOv5x model
model = torch.hub.load('ultralytics/yolov5', 'yolov5x', pretrained=True).to(device)

# Generate distinct colors for each class using HSV color space
def generate_distinct_colors(num_classes):
    colors = {}
    for i, class_name in enumerate(model.names):
        # Use HSV to generate evenly distributed hues
        hue = (i * 255 // num_classes)
        # Convert HSV to BGR (OpenCV uses BGR)
        hsv_color = np.uint8([[[hue, 255, 255]]])
        bgr_color = cv2.cvtColor(hsv_color, cv2.COLOR_HSV2BGR)[0][0]
        # Store as tuple for easier use
        colors[class_name] = tuple(map(int, bgr_color))
    return colors

# Generate colors once at startup
CLASS_COLORS = generate_distinct_colors(len(model.names))

def preprocess_image(image):
    image = Image.fromarray(image)
    image = image.convert("RGB")
    return image

def detect_objects(image):
    image = preprocess_image(image)
    results = model(image)
    image = np.array(image)
    
    # Process all detections at once
    detections = results.xyxy[0]
    for *box, conf, cls in detections:
        x1, y1, x2, y2 = map(int, box)
        class_name = model.names[int(cls)]
        confidence = conf.item() * 100
        color = CLASS_COLORS[class_name]
        
        # Draw rectangle and label
        cv2.rectangle(image, (x1, y1), (x2, y2), color, 4)
        label = f"{class_name} ({confidence:.1f}%)"
        cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, color, 3, cv2.LINE_AA)
    
    return image

# Gradio interface
iface = gr.Interface(
    fn=detect_objects,
    inputs=gr.Image(type="numpy", label="Upload Image"),
    outputs=gr.Image(type="numpy", label="Detected Objects"),
    title="Object Detection with YOLOv5",
    description="Use webcam or upload an image to detect objects.",
    allow_flagging="never",
    examples=["spring_street_after.jpg", "pexels-hikaique-109919.jpg"]
)

iface.launch()