Update app.py

app.py

@@ -42,8 +42,8 @@ FFMPEG_AVAILABLE = check_ffmpeg()
 class BYTETracker:
     def __init__(self, track_thresh=0.3, track_buffer=90, match_thresh=0.5, frame_rate=30):
         self.track_thresh = track_thresh
-        self.track_buffer = track_buffer  # Increased to 90 frames (3 seconds at 30 fps)
-        self.match_thresh = match_thresh  # Lowered to 0.5 for better matching
+        self.track_buffer = track_buffer
+        self.match_thresh = match_thresh
         self.frame_rate = frame_rate
         self.next_id = 1
         self.tracks = {}
@@ -63,7 +63,6 @@ class BYTETracker:
             best_iou = 0
             best_track_id = None
             
-            # Try to match with existing tracks
             for track_id, track_info in self.tracks.items():
                 if current_time - track_info['last_seen'] > self.track_buffer / self.frame_rate:
                     continue
@@ -160,7 +159,7 @@ class BYTETracker:
         iou = intersection_area / (box1_area + box2_area - intersection_area)
         return iou
     
-    def _is_same_worker(self, pos1, pos2, threshold=300):  # Increased to 300 pixels
+    def _is_same_worker(self, pos1, pos2, threshold=300):
         x1, y1 = pos1
         x2, y2 = pos2
         distance = np.sqrt((x1 - x2)**2 + (y1 - y2)**2)
@@ -209,14 +208,15 @@ CONFIG = {
     "VIOLATION_COOLDOWN": 30.0,
     "WORKER_TRACKING_DURATION": 5.0,
     "MAX_PROCESSING_TIME": 60,
-    "FRAME_SKIP": 1,  # Reduced to 1 to process every frame
-    "BATCH_SIZE": 8,
+    "FRAME_SKIP": 1,
+    "BATCH_SIZE": 4,  # Reduced to 4 to lower memory usage
     "PARALLEL_WORKERS": max(1, cpu_count() - 1),
-    "TRACK_BUFFER": 90,  # Increased to 90
+    "TRACK_BUFFER": 90,
     "TRACK_THRESH": 0.3,
-    "MATCH_THRESH": 0.5,  # Lowered to 0.5
+    "MATCH_THRESH": 0.5,
     "SNAPSHOT_QUALITY": 95,
-    "MAX_WORKER_DISTANCE": 300  # Increased to 300
+    "MAX_WORKER_DISTANCE": 300,
+    "TARGET_RESOLUTION": (640, 360)  # Added to resize frames
 }
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -235,6 +235,9 @@ def load_model():
                 torch.hub.download_url_to_file('https://github.com/ultralytics/assets/releases/download/v8.3.0/yolov8n.pt', model_path)
                 
         model = YOLO(model_path).to(device)
+        # Enable FP16 inference if on GPU
+        if device.type == "cuda":
+            model.model.half()
         logger.info(f"Model classes: {model.names}")
         return model
     except Exception as e:
@@ -245,6 +248,9 @@ model = load_model()
 
 # ========================== # Helper Functions # ==========================
 def preprocess_frame(frame):
+    # Resize frame to target resolution
+    target_res = CONFIG["TARGET_RESOLUTION"]
+    frame = cv2.resize(frame, target_res, interpolation=cv2.INTER_LINEAR)
     frame = cv2.convertScaleAbs(frame, alpha=1.2, beta=20)
     return frame
 
@@ -542,9 +548,8 @@ def process_video(video_data, temp_dir):
             frame_rate=fps
         )
 
-        # Modified: Use a single worker ID since we know there's only one worker
-        worker_id_mapping = {}  # Map tracker IDs to a single worker ID
-        unique_violations = {}  # Keyed by violation type to ensure uniqueness
+        worker_id_mapping = {}
+        unique_violations = {}
         snapshots = []
         start_time = time.time()
         frame_skip = CONFIG["FRAME_SKIP"]
@@ -579,17 +584,30 @@ def process_video(video_data, temp_dir):
                 break
 
             try:
-                results = model(batch_frames, device=device, conf=0.1, verbose=False)
+                # Convert frames to tensor and move to device
+                batch_frames_tensor = [torch.from_numpy(frame).permute(2, 0, 1).float() / 255.0 for frame in batch_frames]
+                batch_frames_tensor = torch.stack(batch_frames_tensor).to(device)
+                if device.type == "cuda":
+                    batch_frames_tensor = batch_frames_tensor.half()
+
+                results = model(batch_frames_tensor, device=device, conf=0.1, verbose=False)
             except Exception as e:
                 logger.error(f"Model inference failed: {e}")
                 raise ValueError(f"Failed to process video frames with YOLO model: {str(e)}")
+            finally:
+                # Clear memory
+                batch_frames = []  # Clear the list to free memory
+                if device.type == "cuda":
+                    torch.cuda.empty_cache()
 
             for i, (result, frame_idx) in enumerate(zip(results, batch_indices)):
                 current_time = frame_idx / fps
                 
-                if time.time() - start_time > 1.0:
+                if time.time() - start_time > 0.5:  # Yield more frequently
                     progress = (processed_frames / total_frames) * 100
-                    yield f"Processing video... {progress:.1f}% complete (Frame {processed_frames}/{total_frames})", "", "", "", ""
+                    elapsed_time = time.time() - start_time
+                    fps_processed = processed_frames / elapsed_time if elapsed_time > 0 else 0
+                    yield f"Processing video... {progress:.1f}% complete (Frame {processed_frames}/{total_frames}, {fps_processed:.1f} FPS)", "", "", "", ""
                     start_time = time.time()
 
                 boxes = result.boxes
@@ -631,16 +649,13 @@ def process_video(video_data, temp_dir):
                     if label is None:
                         continue
                     
-                    # Map all tracker IDs to a single worker ID (since we know there's only one worker)
                     if not worker_id_mapping:
-                        worker_id_mapping[tracker_id] = 1  # First worker ID is 1
+                        worker_id_mapping[tracker_id] = 1
                     else:
-                        # Map all subsequent tracker IDs to the same worker ID
                         worker_id_mapping[tracker_id] = worker_id_mapping[list(worker_id_mapping.keys())[0]]
                     
                     worker_id = worker_id_mapping[tracker_id]
                     
-                    # Use violation type as key to ensure uniqueness per worker
                     violation_key = (worker_id, label)
                     
                     if violation_key not in unique_violations:
@@ -654,7 +669,8 @@ def process_video(video_data, temp_dir):
                             "timestamp": current_time
                         }
                         
-                        snapshot_frame = batch_frames[i].copy()
+                        snapshot_frame = batch_frames_tensor[i].cpu().numpy().transpose(1, 2, 0) * 255
+                        snapshot_frame = snapshot_frame.astype(np.uint8)
                         snapshot_frame = draw_detections(snapshot_frame, [detection])
                         
                         cv2.putText(
@@ -687,6 +703,10 @@ def process_video(video_data, temp_dir):
                         
                         logger.info(f"Captured snapshot for {label} violation by worker {worker_id} at {current_time:.2f}s")
 
+                # Clear snapshots periodically to reduce memory usage
+                if len(snapshots) > 100:
+                    snapshots = snapshots[-10:]  # Keep only the last 10 snapshots in memory
+
         cap.release()
         processing_time = time.time() - start_time
         logger.info(f"Processing complete in {processing_time:.2f}s")
@@ -764,6 +784,8 @@ def process_video(video_data, temp_dir):
                 logger.info(f"Cleaned up temporary video file: {video_path}")
             except Exception as e:
                 logger.error(f"Failed to clean up temporary video file {video_path}: {e}")
+        if device.type == "cuda":
+            torch.cuda.empty_cache()
 
 def gradio_interface(video_file):
     temp_dir = None
@@ -808,6 +830,8 @@ def gradio_interface(video_file):
         if temp_dir and os.path.exists(temp_dir):
             shutil.rmtree(temp_dir, ignore_errors=True)
             logger.info(f"Cleaned up temporary directory: {temp_dir}")
+        if device.type == "cuda":
+            torch.cuda.empty_cache()
 
 # ========================== # Gradio Interface # ==========================
 interface = gr.Interface(