Update app.py

app.py

@@ -6,228 +6,118 @@ import time
 import os
 from pathlib import Path
 
-# Create cache directory for models if it doesn't exist
+# Create cache directory for models
 os.makedirs("models", exist_ok=True)
 
-# Check device availability - Hugging Face Spaces often provides GPU
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"Using device: {device}")
 
-# Load YOLOv5x model with caching for faster startup
+# Load YOLOv5x model
 model_path = Path("models/yolov5x.pt")
 if model_path.exists():
     print(f"Loading model from cache: {model_path}")
-    model = torch.hub.load("ultralytics/yolov5", "yolov5x", pretrained=True, 
-                          source="local", path=str(model_path)).to(device)
+    model = torch.hub.load("ultralytics/yolov5", "yolov5x", pretrained=True, source="local", path=str(model_path)).to(device)
 else:
     print("Downloading YOLOv5x model and caching...")
     model = torch.hub.load("ultralytics/yolov5", "yolov5x", pretrained=True).to(device)
-    # Cache the model for faster startup next time
     torch.save(model.state_dict(), model_path)
 
-# 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
+# Model configurations
+model.conf = 0.3  # Confidence threshold
+model.iou = 0.3   # IoU threshold
+model.classes = None  # Detect all classes
 
-# Optimize for GPU if available
 if device.type == "cuda":
-    # Use mixed precision for performance boost
     model.half()
 else:
-    # On CPU, optimize operations
     torch.set_num_threads(os.cpu_count())
 
-# 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
-# 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
+np.random.seed(42)
+colors = np.random.uniform(0, 255, size=(len(model.names), 3))
 
-# Track performance metrics
 total_inference_time = 0
 inference_count = 0
 
 def detect_objects(image):
-    
     global total_inference_time, inference_count
     
     if image is None:
         return None
     
     start_time = time.time()
-    
-    # Create a copy for drawing results
     output_image = image.copy()
-    
-    # Fixed input size for optimal processing
     input_size = 640
     
-    # Perform inference with no gradient calculation
     with torch.no_grad():
-        # Convert image to tensor for faster processing
         results = model(image, size=input_size)
     
-    # Record inference time (model processing only)
     inference_time = time.time() - start_time
     total_inference_time += inference_time
     inference_count += 1
     avg_inference_time = total_inference_time / inference_count
     
-    # Extract detections from first (and only) image
     detections = results.pred[0].cpu().numpy()
     
     for *xyxy, conf, cls in detections:
         x1, y1, x2, y2 = map(int, xyxy)
         class_id = int(cls)
+        color = colors[class_id].tolist()
         
-        color = colors[class_id]
-        
-        cv2.rectangle(output_image, (x1, y1), (x2, y2), color, 3)
+        # Thinner, stylish bounding boxes
+        cv2.rectangle(output_image, (x1, y1), (x2, y2), color, 2, lineType=cv2.LINE_AA)
         
         label = f"{model.names[class_id]} {conf:.2f}"
-        
-        font_scale = 0.7
-        font_thickness = 2
-        
+        font_scale, font_thickness = 0.7, 1
         (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, font_scale, font_thickness)
         
-        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.rectangle(output_image, (x1, y1 - h - 10), (x1 + w + 10, y1), color, -1)
         cv2.putText(output_image, label, (x1 + 5, y1 - 5),
-                    cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), font_thickness)
+                    cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), font_thickness, lineType=cv2.LINE_AA)
     
-    # Calculate FPS
     fps = 1 / inference_time
     
-
-    h, w = output_image.shape[:2]
- 
+    # Stylish FPS display
     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)
-    
-    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)
+    cv2.rectangle(overlay, (10, 10), (300, 80), (0, 0, 0), -1)
+    output_image = cv2.addWeighted(overlay, 0.6, output_image, 0.4, 0)
+    cv2.putText(output_image, f"FPS: {fps:.2f}", (20, 40),
+                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2, lineType=cv2.LINE_AA)
+    cv2.putText(output_image, f"Avg FPS: {1/avg_inference_time:.2f}", (20, 70),
+                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2, lineType=cv2.LINE_AA)
     
     return output_image
 
-# Define example images - these will be stored in the same directory as this script
-example_images = [
-    "spring_street_after.jpg", 
-    "pexels-hikaique-109919.jpg"
-]
-
-# Make sure example directory exists
+example_images = ["spring_street_after.jpg", "pexels-hikaique-109919.jpg"]
 os.makedirs("examples", exist_ok=True)
 
-# Create Gradio interface - optimized for Hugging Face Spaces
 with gr.Blocks(title="Optimized YOLOv5 Object Detection") as demo:
     gr.Markdown("""
     # Optimized YOLOv5 Object Detection
-    
-    This system utilizes YOLOv5 to detect 80+ object types from the COCO dataset.
-    
-    **Performance Features:**
-    - Processing speed: Optimized for 30+ FPS at 640x640 resolution
-    - Confidence threshold: 0.3
-    - IoU threshold: 0.3
-    
-    Upload an image, then click Submit to see the detections!
+    Detects objects using YOLOv5 with enhanced visualization and FPS tracking.
     """)
     
     with gr.Row():
         with gr.Column(scale=1):
             input_image = gr.Image(label="Input Image", type="numpy")
-            with gr.Row():
-                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")
     
+    with gr.Row():
+        clear_button = gr.Button("Clear")
+        submit_button = gr.Button("Submit", variant="primary")
+    
     gr.Examples(
         examples=example_images,
         inputs=input_image,
         outputs=output_image,
         fn=detect_objects,
-        cache_examples=True  # Cache for faster response
+        cache_examples=True
     )
     
     submit_button.click(fn=detect_objects, inputs=input_image, outputs=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
-    )
+    clear_button.click(lambda: (None, None), None, [input_image, output_image])
 
-# Launch for Hugging Face Spaces
-demo.launch()
+demo.launch()