Update app.py

app.py

@@ -43,22 +43,44 @@ 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))
+# Generate more attractive, vibrant colors
+colors = []
+for i in range(len(model.names)):
+    # Use HSV color space for more vibrant colors
+    hue = i / len(model.names)
+    # Full saturation and value for vivid colors
+    saturation = 0.9
+    value = 1.0
+    # Convert HSV to RGB
+    h = hue * 360
+    s = saturation
+    v = value
+    c = v * s
+    x = c * (1 - abs((h / 60) % 2 - 1))
+    m = v - c
+    
+    if h < 60:
+        r, g, b = c, x, 0
+    elif h < 120:
+        r, g, b = x, c, 0
+    elif h < 180:
+        r, g, b = 0, c, x
+    elif h < 240:
+        r, g, b = 0, x, c
+    elif h < 300:
+        r, g, b = x, 0, c
+    else:
+        r, g, b = c, 0, x
+    
+    r, g, b = (r + m) * 255, (g + m) * 255, (b + m) * 255
+    colors.append([int(b), int(g), int(r)])  # OpenCV uses BGR
 
 # Track performance metrics
 total_inference_time = 0
 inference_count = 0
 
 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
-    """
     global total_inference_time, inference_count
     
     if image is None:
@@ -86,53 +108,78 @@ def detect_objects(image):
     # 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()
+        color = colors[class_id]
         
-        cv2.rectangle(output_image, (x1, y1), (x2, y2), color, 4)
+        cv2.rectangle(output_image, (x1, y1), (x2, y2), color, 3)
         
-        # Create label with class name and confidence score
         label = f"{model.names[class_id]} {conf:.2f}"
         
-        font_scale = 0.8
+        font_scale = 0.7
         font_thickness = 2
         
         (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, font_scale, font_thickness)
-  
-        cv2.rectangle(output_image, (x1, y1 - h - 10), (x1 + w + 10, y1), color, -1)
+        
+        alpha = 0.7  
+        overlay = output_image.copy()
+        cv2.rectangle(overlay, (x1, y1 - h - 10), (x1 + w + 10, y1), color, -1)
+        output_image = cv2.addWeighted(overlay, alpha, output_image, 1 - alpha, 0)
         
         cv2.putText(output_image, label, (x1 + 5, y1 - 5),
                     cv2.FONT_HERSHEY_SIMPLEX, font_scale, (0, 0, 0), font_thickness + 1)
+
         cv2.putText(output_image, label, (x1 + 5, y1 - 5),
                     cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), font_thickness)
     
     # Calculate FPS
     fps = 1 / inference_time
     
-    fps_overlay = output_image.copy()
-    cv2.rectangle(fps_overlay, (5, 5), (250, 80), (0, 0, 0), -1)
-    # Apply the overlay with transparency
-    alpha = 0.7
-    output_image = cv2.addWeighted(fps_overlay, alpha, output_image, 1 - alpha, 0)
+
+    h, w = output_image.shape[:2]
+ 
+    overlay = output_image.copy()
+    
+    fps_bg_height = 90
+    fps_bg_width = 200
+    fps_bg_corner = 15  
+    
+    for i in range(fps_bg_height):
+        alpha = 0.8 - (i / fps_bg_height * 0.3) 
+        color_value = int(220 * (1 - i / fps_bg_height)) 
+        cv2.rectangle(overlay, 
+                      (10, 10 + i), 
+                      (fps_bg_width, 10 + i), 
+                      (40, color_value, 40), 
+                      -1)
+
+    cv2.addWeighted(overlay, 0.8, output_image, 0.2, 0, output_image, 
+                    dst=output_image[10:10+fps_bg_height, 10:10+fps_bg_width])
+    
+    cv2.rectangle(output_image, 
+                  (10, 10), 
+                  (fps_bg_width, 10 + fps_bg_height), 
+                  (255, 255, 255), 
+                  2, 
+                  cv2.LINE_AA)
     
-    # Display FPS with larger font
-    cv2.putText(output_image, f"FPS: {fps:.2f}", (10, 35),
-                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
-    cv2.putText(output_image, f"Avg FPS: {1/avg_inference_time:.2f}", (10, 70),
-                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
+    cv2.putText(output_image, "Performance", (20, 35),
+                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2, cv2.LINE_AA)
+    
+    cv2.putText(output_image, f"Current: {fps:.1f} FPS", (20, 65),
+                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2, cv2.LINE_AA)
+    
+    cv2.putText(output_image, f"Average: {1/avg_inference_time:.1f} FPS", (20, 90),
+                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2, cv2.LINE_AA)
     
     return output_image
 
 # Define example images - these will be stored in the same directory as this script
 example_images = [
-    "examples/spring_street_after.jpg", 
-    "examples/pexels-hikaique-109919.jpg"
+    "spring_street_after.jpg", 
+    "pexels-hikaique-109919.jpg"
 ]
 
 # Make sure example directory exists
@@ -150,20 +197,19 @@ with gr.Blocks(title="Optimized YOLOv5 Object Detection") as demo:
     - Confidence threshold: 0.3
     - IoU threshold: 0.3
     
-    Simply upload an image and click Submit to see the detections!
+    Upload an image, then click Submit to see the detections!
     """)
     
     with gr.Row():
         with gr.Column(scale=1):
             input_image = gr.Image(label="Input Image", type="numpy")
             with gr.Row():
-                submit_button = gr.Button("Submit", variant="primary")
-                clear_button = gr.Button("Clear")
+                clear_button = gr.Button("Clear", size="sm")
+                submit_button = gr.Button("Submit", variant="primary", size="lg")
                 
         with gr.Column(scale=1):
             output_image = gr.Image(label="Detected Objects", type="numpy")
     
-    # Example gallery
     gr.Examples(
         examples=example_images,
         inputs=input_image,
@@ -172,10 +218,16 @@ with gr.Blocks(title="Optimized YOLOv5 Object Detection") as demo:
         cache_examples=True  # Cache for faster response
     )
     
-    # Set up button event handlers
     submit_button.click(fn=detect_objects, inputs=input_image, outputs=output_image)
-    clear_button.click(lambda: None, None, [input_image, output_image])
-    
+    clear_button.click(
+        fn=lambda: (None, None),  
+        outputs=[input_image, output_image],
+        queue=False
+    ).then(
+        fn=detect_objects, 
+        inputs=input_image,
+        outputs=output_image
+    )
 
 # Launch for Hugging Face Spaces
 demo.launch()