Update app.py

app.py

@@ -1,5 +1,4 @@
 import torch
-import cv2
 import numpy as np
 import gradio as gr
 from PIL import Image
@@ -7,57 +6,34 @@ from PIL import Image
 # Device configuration
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-# Load YOLOv5x model
-model = torch.hub.load('ultralytics/yolov5', 'yolov5x', pretrained=True).to(device)
+# Load YOLOv5s model (smallest and fastest variant)
+model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True).to(device)
 
-# Generate distinct colors for each class using HSV color space
-def generate_distinct_colors(num_classes):
-    colors = {}
-    for i, class_name in enumerate(model.names):
-        # Use HSV to generate evenly distributed hues
-        hue = (i * 255 // num_classes)
-        # Convert HSV to BGR (OpenCV uses BGR)
-        hsv_color = np.uint8([[[hue, 255, 255]]])
-        bgr_color = cv2.cvtColor(hsv_color, cv2.COLOR_HSV2BGR)[0][0]
-        # Store as tuple for easier use
-        colors[class_name] = tuple(map(int, bgr_color))
-    return colors
-
-# Generate colors once at startup
-CLASS_COLORS = generate_distinct_colors(len(model.names))
-
-def preprocess_image(image):
-    image = Image.fromarray(image)
-    image = image.convert("RGB")
-    return image
+# Enable half-precision for CUDA devices
+if device.type == 'cuda':
+    model.half()
 
 def detect_objects(image):
-    image = preprocess_image(image)
-    results = model(image)
-    image = np.array(image)
+    # Convert numpy array to PIL Image
+    image_pil = Image.fromarray(image)
+    
+    # Perform inference without gradient calculation
+    with torch.no_grad():
+        results = model(image_pil)
     
-    # Process all detections at once
-    detections = results.xyxy[0]
-    for *box, conf, cls in detections:
-        x1, y1, x2, y2 = map(int, box)
-        class_name = model.names[int(cls)]
-        confidence = conf.item() * 100
-        color = CLASS_COLORS[class_name]
-        
-        # Draw rectangle and label
-        cv2.rectangle(image, (x1, y1), (x2, y2), color, 4)
-        label = f"{class_name} ({confidence:.1f}%)"
-        cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, color, 3, cv2.LINE_AA)
+    # Render detections using optimized YOLOv5 method
+    rendered_images = results.render()
     
-    return image
+    # Return the first rendered image as numpy array
+    return np.array(rendered_images[0]) if rendered_images else image
 
 # 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="Object Detection with YOLOv5",
-    description="Use webcam or upload an image to detect objects.",
+    title="High-Speed Object Detection with YOLOv5s",
+    description="Optimized for speed using YOLOv5s and GPU acceleration.",
     allow_flagging="never",
     examples=["spring_street_after.jpg", "pexels-hikaique-109919.jpg"]
 )