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()