Update app.py

app.py

@@ -6,6 +6,8 @@ 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)
@@ -24,12 +26,13 @@ else:
     torch.save(model.state_dict(), model_path)
 
 # Optimize model for speed
-model.conf = 0.3 
-model.iou = 0.3  
-model.classes = None  
+model.conf = 0.25  # Slightly lower confidence threshold
+model.iou = 0.45  # Better IoU threshold
+model.classes = None
+model.max_det = 100  # Limit maximum detections
 
 if device.type == "cuda":
-    model.half()  
+    model.half()  # Use FP16 precision
 else:
     torch.set_num_threads(os.cpu_count())
 
@@ -39,6 +42,27 @@ 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 == "":
@@ -57,52 +81,131 @@ def process_video(video_path):
     frame_height = int(cap.get(4))
     fps = cap.get(cv2.CAP_PROP_FPS)
     
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    # Used h264 codec for better performance
+    fourcc = cv2.VideoWriter_fourcc(*'avc1')
     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()
+    
+    # Start processing thread
+    processing_thread = Thread(target=process_frame, args=(model, frame_queue, result_queue))
+    processing_thread.daemon = True
+    processing_thread.start()
+    
     total_frames = 0
-    total_time = 0
-
-    while cap.isOpened():
-        ret, frame = cap.read()
-        if not ret:
-            break  
-        
-        start_time = time.time()
-        
-        # Convert frame for YOLOv5
-        img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        results = model(img, size=640)
-        
-        inference_time = time.time() - start_time
-        total_time += inference_time
-        total_frames += 1
+    start_time = time.time()
+    processing_started = False
+    frames_buffer = {}
+    next_frame_to_write = 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
+            
+            # 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
+            
+        # Signal thread to finish and process remaining frames
+        frame_queue.put(None)
         
-        detections = results.xyxy[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()
-            cv2.rectangle(frame, (x1, y1), (x2, y2), color, 3, lineType=cv2.LINE_AA)
-            label = f"{model.names[class_id]} {conf:.2f}"
-            cv2.putText(frame, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 255, 255), 2)
-
-        # Calculate FPS
-        avg_fps = total_frames / total_time if total_time > 0 else 0
-        cv2.putText(frame, f"FPS: {avg_fps:.2f}", (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
-
-        out.write(frame)
-
-    cap.release()
-    out.release()
+        # 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)
+            
+            # 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
+    
+    finally:
+        cap.release()
+        out.release()
     
     return output_path
 
 def process_image(image):
+    if image is None:
+        return None
+        
     img = np.array(image)
-    results = model(img, size=640)
+    
+    # Process with smaller size for speed
+    results = model(img, size=512)
 
     detections = results.pred[0].cpu().numpy()
 
@@ -110,9 +213,10 @@ def process_image(image):
         x1, y1, x2, y2 = map(int, xyxy)
         class_id = int(cls)
         color = colors[class_id].tolist()
-        cv2.rectangle(img, (x1, y1), (x2, y2), color, 3, lineType=cv2.LINE_AA)
+        cv2.rectangle(img, (x1, y1), (x2, y2), color, 2, lineType=cv2.LINE_AA)
         label = f"{model.names[class_id]} {conf:.2f}"
-        cv2.putText(img, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 255, 255), 2)
+        # Black text
+        cv2.putText(img, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 0), 2, cv2.LINE_AA)
 
     return Image.fromarray(img)
 
@@ -236,4 +340,5 @@ with gr.Blocks(css=css, title="Video & Image Object Detection by YOLOv5") as dem
     This application enables seamless object detection using the YOLOv5 model, allowing users to analyze images and videos with high accuracy and efficiency.
     """, elem_classes="footer")
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()