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

# Load YOLOv5 model
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = torch.hub.load('ultralytics/yolov5', 'yolov5x', pretrained=True).to(device) 

# Get class names from the model
CLASS_NAMES = model.names  

# Generate consistent colors for each class
random.seed(42)  # Fix the seed for consistent colors
CLASS_COLORS = {cls: (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) for cls in CLASS_NAMES}

def preprocess_image(image):
    """Convert numpy image to PIL format for YOLOv5 processing."""
    image = Image.fromarray(image)
    image = image.convert("RGB")
    return image

def detect_objects(image):
    """Detect objects in the image and draw bounding boxes with consistent colors."""
    image = preprocess_image(image)
    results = model([image])  # YOLOv5 inference
    image = np.array(image)  # Convert PIL image back to numpy for OpenCV

    for *box, conf, cls in results.xyxy[0]:  
        x1, y1, x2, y2 = map(int, box)
        class_name = CLASS_NAMES[int(cls)]  
        confidence = conf.item() * 100  

        color = CLASS_COLORS[class_name]  # Use pre-generated consistent color

        # Draw bounding box
        cv2.rectangle(image, (x1, y1), (x2, y2), color, 4)  

        # Display class label with confidence score
        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


# Create 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"]
)

# Launch the app
iface.launch()