app.py

import torch
import cv2
import numpy as np
import gradio as gr
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)
model.eval()

# Use half-precision if CUDA is available
if device.type == 'cuda':
    model.half()

# Get class names
CLASS_NAMES = model.names  

# Assign random colors for each class
CLASS_COLORS = {cls: (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) for cls in CLASS_NAMES}

def detect_objects(image):
    """Detect objects in an image using YOLOv5 with optimized inference speed."""
    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)  # Convert to BGR for OpenCV
    img_resized = cv2.resize(image, (640, 640))  # Resize for faster processing
    img_tensor = torch.from_numpy(img_resized).to(device).float() / 255.0  # Normalize
    img_tensor = img_tensor.permute(2, 0, 1).unsqueeze(0)  # Convert to batch format

    if device.type == 'cuda':
        img_tensor = img_tensor.half()  # Use half precision for speed

    # Run model inference
    with torch.no_grad():
        results = model(img_tensor)

    detections = results.xyxy[0].cpu().numpy()

    for x1, y1, x2, y2, conf, cls in detections:
        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
        class_name = CLASS_NAMES[int(cls)]
        confidence = conf * 100

        color = CLASS_COLORS[class_name]

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

        # Label
        label = f"{class_name} ({confidence:.1f}%)"
        cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.8, color, 2)

    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # Convert back to RGB for Gradio

# 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="Fast Object Detection with YOLOv5",
    description="Use webcam or upload an image for object detection results.",
    allow_flagging="never"
)

# Launch the app
iface.launch()