app.py

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

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

   
    img_tensor = torch.from_numpy(image).permute(2, 0, 1).float().to(device) / 255.0
    img_tensor = img_tensor.unsqueeze(0)  

    with torch.no_grad():
        results = model(img_tensor)  
    
    detections = results.xyxy[0].cpu().numpy()  

    img_cv = image.copy() 

    for det in detections:
        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="Optimized for 30+ FPS real-time object detection!",
    allow_flagging="never",
)

iface.launch()