Spaces:
Running
Running
File size: 2,167 Bytes
caff61e e82b28e bccf53b dc80d48 248b9ce caff61e f54253a e82b28e f54253a e9a3a08 248b9ce f54253a e9a3a08 36e1064 f54253a a29d5e2 dc80d48 f54253a 73df658 f54253a e82b28e f54253a 248b9ce e82b28e f54253a 73df658 f54253a a29d5e2 f54253a 46e3370 f54253a e82b28e 1195707 f54253a e82b28e 46e3370 f54253a 73df658 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
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()
|