Update app.py

app.py

@@ -3,22 +3,40 @@ import numpy as np
 import gradio as gr
 import cv2
 import time
+import os
+from pathlib import Path
 
-# Check device availability
+# Create cache directory for models if it doesn't exist
+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 (larger model for better accuracy)
-model = torch.hub.load("ultralytics/yolov5", "yolov5x", pretrained=True).to(device)
+# Load YOLOv5x model with caching for faster startup
+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)
+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
 
-# Enable half-precision for GPU acceleration
+# Optimize for GPU if available
 if device.type == "cuda":
-    model.half()  # Use FP16 for performance boost
+    # 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()
@@ -27,6 +45,10 @@ model.eval()
 np.random.seed(42)  # For reproducible colors
 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):
     """
     Process input image for object detection using YOLOv5
@@ -37,22 +59,30 @@ def detect_objects(image):
     Returns:
         output_image: Image with detection results visualized
     """
-    start_time = time.time()
+    global total_inference_time, inference_count
     
-    # 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)
+    if image is None:
+        return None
+    
+    start_time = time.time()
     
     # Create a copy for drawing results
     output_image = image.copy()
     
-    # Resize input to 640x640 for optimal processing speed
+    # 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()
     
@@ -81,34 +111,67 @@ def detect_objects(image):
         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)
+    # Calculate FPS
+    fps = 1 / inference_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}")
+    cv2.putText(output_image, f"Avg FPS: {1/avg_inference_time:.2f}", (10, 70),
+                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
     
     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.
+# 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"
+]
+
+# Make sure example directory exists
+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
-    - Color-coded bounding boxes with confidence scores
-    """,
-    allow_flagging="never",
-    examples=["spring_street_after.jpg", "pexels-hikaique-109919.jpg"],
-)
+    - Real-time FPS display
+    
+    Simply upload an image or take a photo with your camera to see the detections!
+    """)
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            input_image = gr.Image(label="Input Image", type="numpy")
+            with gr.Row():
+                camera_button = gr.Button("Take Photo from Camera")
+                clear_button = gr.Button("Clear")
+                
+        with gr.Column(scale=1):
+            output_image = gr.Image(label="Detected Objects", type="numpy")
+    
+    # Example gallery
+    gr.Examples(
+        examples=example_images,
+        inputs=input_image,
+        outputs=output_image,
+        fn=detect_objects,
+        cache_examples=True  # Cache for faster response
+    )
+    
+    # Set up the inference call
+    input_image.change(fn=detect_objects, inputs=input_image, outputs=output_image)
+    
+    # Event listeners for buttons
+    camera_button.click(lambda: None, None, input_image, js="() => {document.querySelector('button.webcam').click(); return null}")
+    clear_button.click(lambda: None, None, [input_image, output_image])
 
-# Launch the interface
-if __name__ == "__main__":
-    iface.launch()
+# Launch for Hugging Face Spaces
+demo.launch()