app.py

import torch
import numpy as np
import gradio as gr
import cv2
import time
import os
from pathlib import Path

# Create cache directory for models
os.makedirs("models", exist_ok=True)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# Load YOLOv5n model
model_path = Path("models/yolov5n.pt")
if not model_path.exists():
    torch.hub.download_url_to_file("https://github.com/ultralytics/yolov5/releases/download/v6.2/yolov5n.pt", "models/yolov5n.pt")

model = torch.hub.load("ultralytics/yolov5", "custom", path=str(model_path)).to(device)

# Model optimizations
model.conf = 0.5
model.iou = 0.45
if device.type == "cuda":
    model.half()
else:
    model.float()
    torch.set_num_threads(2)

model.eval()

colors = np.random.rand(80, 3) * 255  # COCO classes

def detect_objects(image):
    if image is None:
        return None
    
    start = time.perf_counter()
    
    # Preprocess
    im = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    im = cv2.resize(im, (320, 320))
    tensor = torch.from_numpy(im).to(device)
    tensor = tensor.half() if device.type == "cuda" else tensor.float()
    tensor = tensor.permute(2, 0, 1).unsqueeze(0) / 255
    
    # Inference
    with torch.no_grad():
        pred = model(tensor)[0]
    
    # Post-process
    pred = pred.float() if device.type == "cpu" else pred.half()
    pred = non_max_suppression(pred, model.conf, model.iou)[0]
    
    # Visualization
    output = image.copy()
    if pred is not None:
        pred[:, :4] = scale_coords(tensor.shape[2:], pred[:, :4], image.shape).round()
        for *xyxy, conf, cls in pred:
            x1, y1, x2, y2 = map(int, xyxy)
            cv2.rectangle(output, (x1, y1), (x2, y2), colors[int(cls)].tolist(), 2)
    
    # FPS counter
    fps = 1 / (time.perf_counter() - start)
    cv2.putText(output, f"FPS: {fps:.1f}", (10, 30), 
                cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
    
    return output

with gr.Blocks() as demo:
    gr.Markdown("# Real-Time Object Detection")
    with gr.Row():
        input_img = gr.Image(label="Input", streaming=True)  # Modified webcam handling
        output_img = gr.Image(label="Output")
    input_img.change(detect_objects, input_img, output_img)

demo.launch()