Update app.py

app.py

@@ -3,50 +3,55 @@ import cv2
 import numpy as np
 import gradio as gr
 from PIL import Image
-import random
 
-# Load YOLOv5 model
+# Device configuration
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True).to(device) 
 
-# Get class names from the model
-CLASS_NAMES = model.names  
+# Load YOLOv5x model
+model = torch.hub.load('ultralytics/yolov5', 'yolov5x', pretrained=True).to(device)
 
-# Generate consistent colors for each class
-random.seed(42)  # Fix the seed for consistent colors
-CLASS_COLORS = {cls: (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) for cls in CLASS_NAMES}
+# 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):
-    """Convert numpy image to PIL format for YOLOv5 processing."""
     image = Image.fromarray(image)
     image = image.convert("RGB")
     return image
 
 def detect_objects(image):
-    """Detect objects in the image and draw bounding boxes with consistent colors."""
     image = preprocess_image(image)
-    results = model([image])  # YOLOv5 inference
-    image = np.array(image)  # Convert PIL image back to numpy for OpenCV
-
-    for *box, conf, cls in results.xyxy[0]:  
+    results = model(image)
+    image = np.array(image)
+    
+    # Process all detections at once
+    detections = results.xyxy[0]
+    for *box, conf, cls in detections:
         x1, y1, x2, y2 = map(int, box)
-        class_name = CLASS_NAMES[int(cls)]  
-        confidence = conf.item() * 100  
-
-        color = CLASS_COLORS[class_name]  # Use pre-generated consistent color
-
-        # Draw bounding box
-        cv2.rectangle(image, (x1, y1), (x2, y2), color, 4)  
-
-        # Display class label with confidence score
+        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)  
-
+        cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, color, 3, cv2.LINE_AA)
+    
     return image
 
-
-# Create Gradio Interface
+# Gradio interface
 iface = gr.Interface(
     fn=detect_objects,
     inputs=gr.Image(type="numpy", label="Upload Image"),
@@ -57,5 +62,4 @@ iface = gr.Interface(
     examples=["spring_street_after.jpg", "pexels-hikaique-109919.jpg"]
 )
 
-# Launch the app
-iface.launch()
+iface.launch()