import torch
import numpy as np
import gradio as gr
import cv2
import time
# Check device availability
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Load smaller YOLOv5 model
model = torch.hub.load("ultralytics/yolov5", "yolov5x", pretrained=True).to(device)
# Optimization configurations
model.conf = 0.3  # Confidence threshold
model.iou = 0.3   # NMS IoU threshold
if device.type == "cuda":
    model.half().to(device)  # Use FP16 for performance boost
model.eval()  # Set model to evaluation mode
# Assign fixed colors to each class for bounding boxes
colors = np.random.uniform(0, 255, size=(len(model.names), 3))
def detect_objects(image):
    start_time = time.time()

    # Convert BGR to RGB (if needed, Gradio might already provide RGB)
    # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # Perform inference
    with torch.no_grad():
        results = model(image, size=640)  # Fixed inference size

    # Process results directly on numpy array
    output_image = image.copy()

    # Extract detections
    detections = results.pred[0].cpu().numpy()

    for *xyxy, conf, cls in detections:
        x1, y1, x2, y2 = map(int, xyxy)
        class_id = int(cls)

        # Draw bounding box
        color = colors[class_id].tolist()
        cv2.rectangle(output_image, (x1, y1), (x2, y2), color, 2)

        # Create label
        label = f"{model.names[class_id]} {conf:.2f}"

        # Draw label background
        (w, h), * = cv2.getTextSize(label, cv2.FONT*HERSHEY_SIMPLEX, 0.5, 1)
        cv2.rectangle(output_image, (x1, y1 - 20), (x1 + w, y1), color, -1)

        # Draw label text
        cv2.putText(output_image, label, (x1, y1 - 5),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
    # Calculate FPS
    fps = 1 / (time.time() - start_time)
    print(f"FPS: {fps:.2f}")
    return output_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="Optimized Object Detection with YOLOv5",
    description="Faster detection using YOLOv5s with FP16 and optimized processing",
    allow_flagging="never",
    examples=["spring_street_after.jpg", "pexels-hikaique-109919.jpg"],
)
iface.launch()