Spaces:
Running
Running
| import torch | |
| import numpy as np | |
| import gradio as gr | |
| import cv2 | |
| import time | |
| # Check device availability | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| print(f"Using device: {device}") | |
| # Load YOLOv5x model (larger model for better accuracy) | |
| model = torch.hub.load("ultralytics/yolov5", "yolov5x", pretrained=True).to(device) | |
| # Optimization configurations | |
| model.conf = 0.3 # Confidence threshold of 0.3 as specified | |
| model.iou = 0.3 # NMS IoU threshold of 0.3 as specified | |
| model.classes = None # Detect all 80+ COCO classes | |
| # Enable half-precision for GPU acceleration | |
| if device.type == "cuda": | |
| model.half() # Use FP16 for performance boost | |
| # Set model to evaluation mode for inference | |
| model.eval() | |
| # Assign fixed colors to each class for consistent visualization | |
| np.random.seed(42) # For reproducible colors | |
| colors = np.random.uniform(0, 255, size=(len(model.names), 3)) | |
| def detect_objects(image): | |
| """ | |
| Process input image for object detection using YOLOv5 | |
| Args: | |
| image: Input image as numpy array | |
| Returns: | |
| output_image: Image with detection results visualized | |
| """ | |
| start_time = time.time() | |
| # Convert image to RGB if it's in BGR format | |
| if image.shape[2] == 3 and image[0,0,0] == image[0,0,2]: | |
| image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| # Create a copy for drawing results | |
| output_image = image.copy() | |
| # Resize input to 640x640 for optimal processing speed | |
| input_size = 640 | |
| # Perform inference with no gradient calculation | |
| with torch.no_grad(): | |
| results = model(image, size=input_size) | |
| # Extract detections from first (and only) image | |
| detections = results.pred[0].cpu().numpy() | |
| # Draw each detection on the output image | |
| for *xyxy, conf, cls in detections: | |
| # Extract coordinates and convert to integers | |
| x1, y1, x2, y2 = map(int, xyxy) | |
| class_id = int(cls) | |
| # Get color for this class | |
| color = colors[class_id].tolist() | |
| # Draw bounding box | |
| cv2.rectangle(output_image, (x1, y1), (x2, y2), color, 2) | |
| # Create label with class name and confidence score | |
| label = f"{model.names[class_id]} {conf:.2f}" | |
| # Calculate text size for background rectangle | |
| (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1) | |
| # Draw label background | |
| cv2.rectangle(output_image, (x1, y1 - 20), (x1 + w, y1), color, -1) | |
| # Draw label text | |
| cv2.putText(output_image, label, (x1, y1 - 5), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1) | |
| # Calculate and display FPS | |
| fps = 1 / (time.time() - start_time) | |
| # Add FPS counter to the image | |
| cv2.putText(output_image, f"FPS: {fps:.2f}", (10, 30), | |
| cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2) | |
| print(f"Detection complete - FPS: {fps:.2f}") | |
| return output_image | |
| # Create 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="Optimized Object Detection with YOLOv5x", | |
| description=""" | |
| This system utilizes YOLOv5x to detect 80+ object types from the COCO dataset. | |
| - Processing speed: Optimized for 30+ FPS at 640x640 resolution | |
| - Confidence threshold: 0.3 | |
| - IoU threshold: 0.3 | |
| - Color-coded bounding boxes with confidence scores | |
| """, | |
| allow_flagging="never", | |
| examples=["spring_street_after.jpg", "pexels-hikaique-109919.jpg"], | |
| ) | |
| # Launch the interface | |
| if __name__ == "__main__": | |
| iface.launch() |