Update app.py

app.py

@@ -6,8 +6,6 @@ import time
 import os
 from pathlib import Path
 from PIL import Image
-from threading import Thread
-from queue import Queue
 
 # Create cache directory for models
 os.makedirs("models", exist_ok=True)
@@ -26,9 +24,9 @@ else:
     torch.save(model.state_dict(), model_path)
 
 # Optimize model for speed
-model.conf = 0.25  # Slightly lower confidence threshold
-model.iou = 0.45  # Better IoU threshold
-model.classes = None
+model.conf = 0.25  # Lower confidence threshold for speed
+model.iou = 0.45   # Better IoU threshold
+model.classes = None  
 model.max_det = 100  # Limit maximum detections
 
 if device.type == "cuda":
@@ -42,27 +40,6 @@ model.eval()
 np.random.seed(42)
 colors = np.random.randint(0, 255, size=(len(model.names), 3), dtype=np.uint8)
 
-# Async video processing
-def process_frame(model, frame_queue, result_queue):
-    while True:
-        if frame_queue.empty():
-            time.sleep(0.001)
-            continue
-            
-        frame_data = frame_queue.get()
-        if frame_data is None:  # Signal to stop
-            result_queue.put(None)
-            break
-            
-        frame, frame_index = frame_data
-        img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        
-        # Use a smaller inference size for speed
-        results = model(img, size=384)  # Reduced from 640 to 384
-        
-        detections = results.xyxy[0].cpu().numpy()
-        result_queue.put((frame, detections, frame_index))
-
 def process_video(video_path):
     # Check if video_path is None or empty
     if video_path is None or video_path == "":
@@ -81,120 +58,68 @@ def process_video(video_path):
     frame_height = int(cap.get(4))
     fps = cap.get(cv2.CAP_PROP_FPS)
     
-    # Used h264 codec for better performance
-    fourcc = cv2.VideoWriter_fourcc(*'avc1')
+    # Use mp4v codec which is more widely supported
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
     output_path = "output_video.mp4"
     out = cv2.VideoWriter(output_path, fourcc, fps, (frame_width, frame_height))
 
-    # Created queues for async processing
-    frame_queue = Queue(maxsize=10)
-    result_queue = Queue()
+    # For FPS calculation
+    frame_count = 0
+    start_time = time.time()
     
-    # Start processing thread
-    processing_thread = Thread(target=process_frame, args=(model, frame_queue, result_queue))
-    processing_thread.daemon = True
-    processing_thread.start()
+    # Skip frames for faster processing if needed
+    frame_skip = 0
+    if device.type != "cuda":  # Skip more frames on CPU
+        frame_skip = 1
     
-    total_frames = 0
-    start_time = time.time()
-    processing_started = False
-    frames_buffer = {}
-    next_frame_to_write = 0
+    frame_idx = 0
     
-    try:
-        while cap.isOpened():
-            ret, frame = cap.read()
-            if not ret:
-                break
-                
-            if not processing_started:
-                processing_started = True
-                start_time = time.time()
-                
-            frame_queue.put((frame, total_frames))
-            total_frames += 1
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
             
-            # Process results if available
-            while not result_queue.empty():
-                result = result_queue.get()
-                if result is None:
-                    break
-                    
-                processed_frame, detections, frame_idx = result
-                frames_buffer[frame_idx] = (processed_frame, detections)
-                
-                # Write frames in order
-                while next_frame_to_write in frames_buffer:
-                    buffer_frame, buffer_detections = frames_buffer.pop(next_frame_to_write)
-                    
-                    # Draw bounding boxes
-                    for *xyxy, conf, cls in buffer_detections:
-                        if conf < 0.35:  # Additional filtering
-                            continue
-                        x1, y1, x2, y2 = map(int, xyxy)
-                        class_id = int(cls)
-                        color = colors[class_id].tolist()
-                        cv2.rectangle(buffer_frame, (x1, y1), (x2, y2), color, 2, lineType=cv2.LINE_AA)
-                        label = f"{model.names[class_id]} {conf:.2f}"
-                        # Black text with white outline for better visibility
-                        cv2.putText(buffer_frame, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 
-                                   0.7, (0, 0, 0), 2, cv2.LINE_AA)
-                    
-                    # Calculate elapsed time and FPS
-                    elapsed = time.time() - start_time
-                    current_fps = next_frame_to_write / elapsed if elapsed > 0 else 0
-                    
-                    # Add FPS counter with black text
-                    cv2.putText(buffer_frame, f"FPS: {current_fps:.2f}", (20, 40), 
-                               cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2, cv2.LINE_AA)
-                    
-                    out.write(buffer_frame)
-                    next_frame_to_write += 1
+        frame_idx += 1
+        if frame_skip > 0 and frame_idx % (frame_skip + 1) != 0:
+            out.write(frame)  # Write original frame
+            continue
+        
+        # Convert frame for YOLOv5
+        img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        
+        # Use smaller inference size for speed
+        results = model(img, size=384)  # Reduced from 640 to 384
+        
+        detections = results.xyxy[0].cpu().numpy()
+        
+        # Draw bounding boxes
+        for *xyxy, conf, cls in detections:
+            x1, y1, x2, y2 = map(int, xyxy)
+            class_id = int(cls)
+            color = colors[class_id].tolist()
+            cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2, lineType=cv2.LINE_AA)
+            label = f"{model.names[class_id]} {conf:.2f}"
             
-        # Signal thread to finish and process remaining frames
-        frame_queue.put(None)
+            # Black text
+            cv2.putText(frame, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 
+                       0.7, (0, 0, 0), 2, cv2.LINE_AA)
         
-        # Process remaining buffered frames
-        while True:
-            if result_queue.empty():
-                time.sleep(0.01)
-                continue
-                
-            result = result_queue.get()
-            if result is None:
-                break
-                
-            processed_frame, detections, frame_idx = result
-            frames_buffer[frame_idx] = (processed_frame, detections)
+        # Update frame count for FPS calculation
+        frame_count += 1
+        
+        # Calculate and display FPS every 10 frames
+        if frame_count % 10 == 0:
+            elapsed_time = time.time() - start_time
+            fps_calc = frame_count / elapsed_time if elapsed_time > 0 else 0
             
-            # Write remaining frames in order
-            while next_frame_to_write in frames_buffer:
-                buffer_frame, buffer_detections = frames_buffer.pop(next_frame_to_write)
-                
-                # Draw bounding boxes
-                for *xyxy, conf, cls in buffer_detections:
-                    if conf < 0.35:
-                        continue
-                    x1, y1, x2, y2 = map(int, xyxy)
-                    class_id = int(cls)
-                    color = colors[class_id].tolist()
-                    cv2.rectangle(buffer_frame, (x1, y1), (x2, y2), color, 2, lineType=cv2.LINE_AA)
-                    label = f"{model.names[class_id]} {conf:.2f}"
-                    cv2.putText(buffer_frame, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 
-                               0.7, (0, 0, 0), 2, cv2.LINE_AA)
-                
-                # Add FPS counter
-                elapsed = time.time() - start_time
-                current_fps = next_frame_to_write / elapsed if elapsed > 0 else 0
-                cv2.putText(buffer_frame, f"FPS: {current_fps:.2f}", (20, 40), 
-                           cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2, cv2.LINE_AA)
-                
-                out.write(buffer_frame)
-                next_frame_to_write += 1
+            # Add FPS counter with black text
+            cv2.putText(frame, f"FPS: {fps_calc:.2f}", (20, 40), 
+                       cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2, cv2.LINE_AA)
+        
+        out.write(frame)
     
-    finally:
-        cap.release()
-        out.release()
+    cap.release()
+    out.release()
     
     return output_path