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

# Check device availability
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

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

if device.type == "cuda":
    model.half()  # Use FP16 for performance boost

# Print available object classes
print(f"Model loaded with {len(model.names)} classes: {model.names}")

# Assign random colors to each class for bounding boxes
colors = {i: [int(c) for c in np.random.randint(0, 255, 3)] for i in range(len(model.names))}

def detect_objects(image):
    start_time = time.time()  # Start FPS measurement

    image_pil = Image.fromarray(image)

    with torch.no_grad():
        results = model(image_pil, conf=0.3, iou=0.3)  # Apply NMS with IoU = 0.3

    rendered_images = results.render()  # Get rendered image with default YOLOv5 visualization

    # Get bounding boxes and draw color-coded boxes
    img_cv = np.array(rendered_images[0]) if rendered_images else image

    for det in results.xyxy[0]:  # Bounding box format: x1, y1, x2, y2, conf, cls
        x1, y1, x2, y2, conf, cls = map(int, det[:6])
        label = f"{model.names[cls]}: {conf:.2f}"

        cv2.rectangle(img_cv, (x1, y1), (x2, y2), colors[cls], 2)
        cv2.putText(img_cv, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, colors[cls], 2)

    # FPS Calculation
    end_time = time.time()
    fps = 1 / (end_time - start_time)
    print(f"FPS: {fps:.2f}")

    return img_cv

# 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. Optimized for speed and accuracy!",
    allow_flagging="never",
    examples=["spring_street_after.jpg", "pexels-hikaique-109919.jpg"],
)

iface.launch()