Update app.py

app.py

@@ -14,18 +14,19 @@ os.makedirs("models", exist_ok=True)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"Using device: {device}")
 
-model_path = Path("models/yolov5x.pt")
+# Use YOLOv5n (nano) for higher FPS
+model_path = Path("models/yolov5n.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", "yolov5n", 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)
+    print("Downloading YOLOv5n model and caching...")
+    model = torch.hub.load("ultralytics/yolov5", "yolov5n", pretrained=True).to(device)
     torch.save(model.state_dict(), model_path)
 
 # Model configurations for better performance
-model.conf = 0.5  # Slightly lower confidence threshold for real-time
-model.iou = 0.45  # Slightly lower IOU threshold for real-time
+model.conf = 0.5  # Confidence threshold
+model.iou = 0.45  # IOU threshold
 model.classes = None  # Detect all classes
 model.max_det = 20  # Limit detections for speed
 
@@ -44,6 +45,8 @@ colors = np.random.uniform(0, 255, size=(len(model.names), 3))
 total_inference_time = 0
 inference_count = 0
 fps_queue = Queue(maxsize=30)  # Store last 30 FPS values for smoothing
+for _ in range(30):  # Initialize with reasonable values
+    fps_queue.put(30.0)
 
 # Threading variables
 processing_lock = threading.Lock()
@@ -107,58 +110,73 @@ def detect_objects(image):
 def process_frame_thread():
     """Background thread for processing frames"""
     while not stop_event.is_set():
-        if not frame_queue.empty():
-            frame = frame_queue.get()
-            
-            # Skip if there's a processing lock (from image upload)
-            if processing_lock.locked():
-                result_queue.put(frame)  # Return unprocessed frame
-                continue
+        try:
+            if not frame_queue.empty():
+                frame = frame_queue.get()
                 
-            # Process the frame
-            with torch.no_grad():  # Ensure no gradients for inference
-                input_size = 384  # Smaller size for real-time processing
-                results = model(frame, size=input_size)
-            
-            # Calculate FPS
-            inference_time = time.time() - frame.get('timestamp', time.time())
-            current_fps = 1 / inference_time if inference_time > 0 else 30
-            
-            # Update rolling FPS average
-            fps_queue.put(current_fps)
-            avg_fps = sum(list(fps_queue.queue)) / fps_queue.qsize()
-            
-            # Draw detections
-            output = frame['image'].copy()
-            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()
+                # Skip if there's a processing lock (from image upload)
+                if processing_lock.locked():
+                    result_queue.put(frame)  # Return unprocessed frame
+                    continue
+                    
+                # Process the frame
+                start_time = time.time()
+                with torch.no_grad():  # Ensure no gradients for inference
+                    input_size = 384  # Smaller size for real-time processing
+                    results = model(frame['image'], size=input_size)
                 
-                # Draw rectangle and label
-                cv2.rectangle(output, (x1, y1), (x2, y2), color, 2, lineType=cv2.LINE_AA)
+                # Calculate FPS
+                inference_time = time.time() - start_time
+                current_fps = 1 / inference_time if inference_time > 0 else 30
                 
-                label = f"{model.names[class_id]} {conf:.2f}"
-                font_scale, font_thickness = 0.6, 1  # Smaller for real-time
-                (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, font_scale, font_thickness)
+                # Update rolling FPS average
+                if not fps_queue.full():
+                    fps_queue.put(current_fps)
+                else:
+                    try:
+                        fps_queue.get_nowait()
+                        fps_queue.put(current_fps)
+                    except:
+                        pass
                 
-                cv2.rectangle(output, (x1, y1 - h - 5), (x1 + w + 5, y1), color, -1)
-                cv2.putText(output, label, (x1 + 3, y1 - 3),
-                            cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), font_thickness, lineType=cv2.LINE_AA)
-            
-            # Add FPS counter
-            cv2.rectangle(output, (10, 10), (210, 80), (0, 0, 0), -1)
-            cv2.putText(output, f"FPS: {current_fps:.1f}", (20, 40),
-                        cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2, lineType=cv2.LINE_AA)
-            cv2.putText(output, f"Avg FPS: {avg_fps:.1f}", (20, 70),
-                        cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2, lineType=cv2.LINE_AA)
-            
-            # Put the processed frame in the result queue
-            result_queue.put({'image': output, 'fps': current_fps})
-        else:
-            time.sleep(0.001)  # Small sleep to prevent CPU spinning
+                fps_values = list(fps_queue.queue)
+                avg_fps = sum(fps_values) / len(fps_values) if fps_values else 30.0
+                
+                # Draw detections
+                output = frame['image'].copy()
+                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()
+                    
+                    # Draw rectangle and label
+                    cv2.rectangle(output, (x1, y1), (x2, y2), color, 2, lineType=cv2.LINE_AA)
+                    
+                    label = f"{model.names[class_id]} {conf:.2f}"
+                    font_scale, font_thickness = 0.6, 1  # Smaller for real-time
+                    (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, font_scale, font_thickness)
+                    
+                    cv2.rectangle(output, (x1, y1 - h - 5), (x1 + w + 5, y1), color, -1)
+                    cv2.putText(output, label, (x1 + 3, y1 - 3),
+                                cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), font_thickness, lineType=cv2.LINE_AA)
+                
+                # Add FPS counter
+                cv2.rectangle(output, (10, 10), (210, 80), (0, 0, 0), -1)
+                cv2.putText(output, f"FPS: {current_fps:.1f}", (20, 40),
+                            cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2, lineType=cv2.LINE_AA)
+                cv2.putText(output, f"Avg FPS: {avg_fps:.1f}", (20, 70),
+                            cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2, lineType=cv2.LINE_AA)
+                
+                # Put the processed frame in the result queue
+                if not result_queue.full():
+                    result_queue.put({'image': output, 'fps': current_fps})
+            else:
+                time.sleep(0.001)  # Small sleep to prevent CPU spinning
+        except Exception as e:
+            print(f"Error in frame processing thread: {e}")
+            time.sleep(0.1)  # Pause briefly on error
 
 def webcam_feed():
     """Generator function for webcam feed"""
@@ -170,20 +188,33 @@ def webcam_feed():
     
     # Open webcam
     cap = cv2.VideoCapture(0)
+    if not cap.isOpened():
+        print("Warning: Unable to open webcam! Using dummy frames instead.")
+        # Create a dummy frame with a message
+        dummy_frame = np.zeros((480, 640, 3), dtype=np.uint8)
+        cv2.putText(dummy_frame, "Webcam not available", (100, 240), 
+                    cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
+        while True:
+            yield dummy_frame
+            time.sleep(0.033)  # ~30 FPS
+    
+    # Set webcam properties for best performance
     cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
     cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
+    cap.set(cv2.CAP_PROP_FPS, 30)  # Request 30 FPS from camera if supported
     
     try:
-        while True:
+        while cap.isOpened():
             success, frame = cap.read()
             if not success:
+                print("Failed to read from webcam")
                 break
                 
-            # Put frame in queue for processing
+            # Put frame in queue for processing if not full
             if not frame_queue.full():
                 frame_queue.put({'image': frame, 'timestamp': time.time()})
             
-            # Get processed frame from result queue
+            # Get processed frame from result queue if available
             if not result_queue.empty():
                 result = result_queue.get()
                 yield result['image']
@@ -191,8 +222,8 @@ def webcam_feed():
                 # If no processed frame is available, yield the raw frame
                 yield frame
                 
-            # Control frame rate to not overwhelm the system
-            time.sleep(0.01)
+            # Control frame rate
+            time.sleep(0.01)  # Small delay to prevent overwhelming the system
     finally:
         cap.release()
 
@@ -245,24 +276,8 @@ with gr.Blocks(title="YOLOv5 Object Detection - Real-time & Image Upload") as de
     submit_button.click(fn=process_uploaded_image, inputs=input_image, outputs=output_image)
     clear_button.click(lambda: (None, None), None, [input_image, output_image])
     
-    # Connect webcam feed
-    demo.load(lambda: None, None, webcam_output, _js="""
-        () => {
-            // Keep the webcam tab refreshing at high frequency
-            setInterval(() => {
-                if (document.querySelector('.tabitem:first-child').style.display !== 'none') {
-                    const webcamImg = document.querySelector('.tabitem:first-child img');
-                    if (webcamImg) {
-                        const src = webcamImg.src;
-                        webcamImg.src = src.includes('?') ? src.split('?')[0] + '?t=' + Date.now() : src + '?t=' + Date.now();
-                    }
-                }
-            }, 33); // ~30 FPS refresh rate
-            return [];
-        }
-    """)
-    
     # Start webcam feed
+    demo.load(fn=lambda: None, inputs=None, outputs=webcam_output)
     webcam_output.update(webcam_feed)
 
 # Cleanup function to stop threads when app closes
@@ -270,5 +285,6 @@ def cleanup():
     stop_event.set()
     print("Cleaning up threads...")
 
+# Register cleanup handler
 demo.close = cleanup
-demo.launch()
+demo.launch(share=False)