Update app.py

app.py

@@ -14,8 +14,6 @@ import base64
 import logging
 from retrying import retry
 import uuid
-from multiprocessing import Pool, cpu_count
-from functools import partial
 
 # ==========================
 # Optimized Configuration
@@ -53,17 +51,17 @@ CONFIG = {
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo2/resolve/main/static/output/",
     "CONFIDENCE_THRESHOLDS": {
-        "no_helmet": 0.6,  # Higher threshold to reduce false positives
+        "no_helmet": 0.6,
         "no_harness": 0.4,
         "unsafe_posture": 0.4,
         "unsafe_zone": 0.4,
         "improper_tool_use": 0.4
     },
-    "MIN_VIOLATION_FRAMES": 3,  # Require 3+ detections to confirm violation
-    "WORKER_TRACKING_DURATION": 3.0,  # Track workers for 3 seconds
-    "MAX_PROCESSING_TIME": 30,  # 30-second limit
-    "PARALLEL_WORKERS": max(1, cpu_count() - 1),  # Use all CPU cores
-    "CHUNK_SIZE": 8  # Frames per batch
+    "MIN_VIOLATION_FRAMES": 3,
+    "WORKER_TRACKING_DURATION": 3.0,
+    "MAX_PROCESSING_TIME": 60,  # 1 minute limit
+    "FRAME_SKIP": 2,  # Process every 2nd frame for speed
+    "BATCH_SIZE": 16  # Frames per batch
 }
 
 # Setup logging
@@ -119,7 +117,6 @@ def calculate_iou(box1, box2):
     x1, y1, w1, h1 = box1
     x2, y2, w2, h2 = box2
     
-    # Calculate intersection area
     x_left = max(x1 - w1/2, x2 - w2/2)
     y_top = max(y1 - h1/2, y2 - h2/2)
     x_right = min(x1 + w1/2, x2 + w2/2)
@@ -135,36 +132,6 @@ def calculate_iou(box1, box2):
     
     return intersection_area / union_area
 
-def process_frame_batch(frame_batch, frame_indices, fps):
-    batch_results = []
-    results = model(frame_batch, device=device, conf=0.1, verbose=False)
-    
-    for idx, (result, frame_idx) in enumerate(zip(results, frame_indices)):
-        current_time = frame_idx / fps
-        detections = []
-        
-        boxes = result.boxes
-        for box in boxes:
-            cls = int(box.cls)
-            conf = float(box.conf)
-            label = CONFIG["VIOLATION_LABELS"].get(cls, None)
-            
-            if label is None or conf < CONFIG["CONFIDENCE_THRESHOLDS"].get(label, 0.25):
-                continue
-
-            bbox = [round(x, 2) for x in box.xywh.cpu().numpy()[0]]
-            detections.append({
-                "frame": frame_idx,
-                "violation": label,
-                "confidence": round(conf, 2),
-                "bounding_box": bbox,
-                "timestamp": current_time
-            })
-        
-        batch_results.append((frame_idx, detections))
-    
-    return batch_results
-
 def generate_violation_pdf(violations, score):
     try:
         pdf_filename = f"violations_{int(time.time())}.pdf"
@@ -227,7 +194,7 @@ def calculate_safety_score(violations):
     return max(score, 0)
 
 # ==========================
-# Optimized Video Processing
+# Fast Video Processing
 # ==========================
 def process_video(video_data):
     try:
@@ -246,62 +213,97 @@ def process_video(video_data):
         total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
         fps = cap.get(cv2.CAP_PROP_FPS)
         if fps <= 0:
-            fps = 30  # Default assumption if FPS not available
+            fps = 30
         duration = total_frames / fps
         width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
         height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
 
         logger.info(f"Video properties: {duration:.2f}s, {total_frames} frames, {fps:.1f} FPS, {width}x{height}")
 
-        # Read all frames upfront
-        all_frames = []
-        all_indices = []
-        for frame_idx in range(total_frames):
-            ret, frame = cap.read()
-            if not ret:
-                break
-            all_frames.append(frame)
-            all_indices.append(frame_idx)
-        cap.release()
-
-        # Process frames in parallel batches
         workers = []
         violations = []
         helmet_violations = {}
         snapshots = []
         start_time = time.time()
+        processed_frames = 0
+        frame_skip = CONFIG["FRAME_SKIP"]
 
-        # Split frames into batches
-        frame_batches = [all_frames[i:i + CONFIG["CHUNK_SIZE"]] for i in range(0, len(all_frames), CONFIG["CHUNK_SIZE"])]
-        frame_indices_batches = [all_indices[i:i + CONFIG["CHUNK_SIZE"]] for i in range(0, len(all_indices), CONFIG["CHUNK_SIZE"])]
-
-        # Process batches in parallel
-        with Pool(processes=CONFIG["PARALLEL_WORKERS"]) as pool:
-            process_func = partial(process_frame_batch, fps=fps)
-            results = pool.starmap(process_func, zip(frame_batches, frame_indices_batches))
+        # Process frames in batches
+        while True:
+            batch_frames = []
+            batch_indices = []
+            
+            # Collect frames for this batch
+            for _ in range(CONFIG["BATCH_SIZE"]):
+                frame_idx = int(cap.get(cv2.CAP_PROP_POS_FRAMES))
+                if frame_idx >= total_frames:
+                    break
+                
+                ret, frame = cap.read()
+                if not ret:
+                    break
+                
+                # Skip frames if needed
+                for _ in range(frame_skip - 1):
+                    if not cap.grab():
+                        break
+                
+                batch_frames.append(frame)
+                batch_indices.append(frame_idx)
+                processed_frames += 1
+
+            # Break if no more frames
+            if not batch_frames:
+                break
 
-        # Flatten results and track workers
-        for batch_result in results:
-            for frame_idx, detections in batch_result:
+            # Run batch detection
+            results = model(batch_frames, device=device, conf=0.1, verbose=False)
+            
+            # Process results for each frame in batch
+            for i, (result, frame_idx) in enumerate(zip(results, batch_indices)):
                 current_time = frame_idx / fps
+                
+                # Update progress periodically
+                if time.time() - start_time > 1.0:  # Update every second
+                    progress = (frame_idx / total_frames) * 100
+                    yield f"Processing video... {progress:.1f}% complete (Frame {frame_idx}/{total_frames})", "", "", "", ""
+                    start_time = time.time()
+
+                # Process detections in this frame
+                boxes = result.boxes
+                for box in boxes:
+                    cls = int(box.cls)
+                    conf = float(box.conf)
+                    label = CONFIG["VIOLATION_LABELS"].get(cls, None)
+                    
+                    if label is None or conf < CONFIG["CONFIDENCE_THRESHOLDS"].get(label, 0.25):
+                        continue
+
+                    bbox = [round(x, 2) for x in box.xywh.cpu().numpy()[0]]
+                    detection = {
+                        "frame": frame_idx,
+                        "violation": label,
+                        "confidence": round(conf, 2),
+                        "bounding_box": bbox,
+                        "timestamp": current_time
+                    }
 
-                for detection in detections:
                     # Worker tracking
                     worker_id = None
                     max_iou = 0
                     for idx, worker in enumerate(workers):
-                        iou = calculate_iou(detection["bounding_box"], worker["bbox"])
+                        iou = calculate_iou(bbox, worker["bbox"])
                         if iou > max_iou and iou > 0.4:  # IOU threshold
                             max_iou = iou
                             worker_id = worker["id"]
-                            workers[idx]["bbox"] = detection["bounding_box"]
+                            workers[idx]["bbox"] = bbox
                             workers[idx]["last_seen"] = current_time
 
                     if worker_id is None:
                         worker_id = len(workers) + 1
                         workers.append({
                             "id": worker_id,
-                            "bbox": detection["bounding_box"],
+                            "bbox": bbox,
                             "first_seen": current_time,
                             "last_seen": current_time
                         })
@@ -319,6 +321,11 @@ def process_video(video_data):
                 # Remove inactive workers
                 workers = [w for w in workers if current_time - w["last_seen"] < CONFIG["WORKER_TRACKING_DURATION"]]
 
+        cap.release()
+        os.remove(video_path)
+        processing_time = time.time() - start_time
+        logger.info(f"Processing complete in {processing_time:.2f}s. {len(violations)} violations found.")
+
         # Confirm helmet violations (require multiple detections)
         for worker_id, detections in helmet_violations.items():
             if len(detections) >= CONFIG["MIN_VIOLATION_FRAMES"]:
@@ -342,10 +349,6 @@ def process_video(video_data):
                     })
                 cap.release()
 
-        os.remove(video_path)
-        processing_time = time.time() - start_time
-        logger.info(f"Processing complete in {processing_time:.2f}s. {len(violations)} violations found.")
-
         # Generate results
         if not violations:
             yield "No violations detected in the video.", "Safety Score: 100%", "No snapshots captured.", "N/A", "N/A"