Update app.py

app.py

@@ -14,9 +14,11 @@ import base64
 import logging
 from retrying import retry
 import uuid
+from multiprocessing import Pool, cpu_count
+from functools import partial
 
 # ==========================
-# Enhanced Configuration
+# Ultra-Fast Configuration
 # ==========================
 CONFIG = {
     "MODEL_PATH": "yolov8_safety.pt",
@@ -50,7 +52,6 @@ CONFIG = {
         "domain": "login"
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo2/resolve/main/static/output/",
-    "FRAME_SKIP": 2,  # Process every 2nd frame for balance of speed/accuracy
     "CONFIDENCE_THRESHOLDS": {
         "no_helmet": 0.6,
         "no_harness": 0.15,
@@ -59,12 +60,12 @@ CONFIG = {
         "improper_tool_use": 0.15
     },
     "IOU_THRESHOLD": 0.4,
-    "MIN_VIOLATION_FRAMES": 3,  # Require more consistent detections
+    "MIN_VIOLATION_FRAMES": 3,
     "HELMET_CONFIDENCE_THRESHOLD": 0.65,
-    "WORKER_TRACKING_DURATION": 3.0,  # Seconds to track a worker
-    "MIN_FRAME_RATE": 5,  # Minimum frames per second to process
-    "MAX_FRAME_RATE": 15,  # Maximum frames per second to process
-    "BATCH_SIZE": 8  # Number of frames to process at once
+    "WORKER_TRACKING_DURATION": 3.0,
+    "MAX_PROCESSING_TIME": 30,  # 30 second hard limit
+    "PARALLEL_WORKERS": max(1, cpu_count() - 1),  # Use all but one CPU core
+    "CHUNK_SIZE": 10  # Frames per parallel batch
 }
 
 # Setup logging
@@ -83,7 +84,7 @@ def load_model():
             logger.info(f"Model loaded: {model_path}")
         else:
             model_path = CONFIG["FALLBACK_MODEL"]
-            logger.warning("Using fallback model. Detection accuracy may be poor. Train yolov8_safety.pt for best results.")
+            logger.warning("Using fallback model. Train yolov8_safety.pt for best results.")
             if not os.path.isfile(model_path):
                 logger.info(f"Downloading fallback model: {model_path}")
                 torch.hub.download_url_to_file('https://github.com/ultralytics/assets/releases/download/v8.3.0/yolov8n.pt', model_path)
@@ -121,7 +122,7 @@ def calculate_iou(box1, box2):
     x1, y1, w1, h1 = box1
     x2, y2, w2, h2 = box2
     
-    # Calculate coordinates of the intersection rectangle
+    # Calculate intersection coordinates
     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)
@@ -137,6 +138,36 @@ 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, iou=CONFIG["IOU_THRESHOLD"], 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"
@@ -199,7 +230,7 @@ def calculate_safety_score(violations):
     return max(score, 0)
 
 # ==========================
-# Optimized Video Processing
+# Ultra-Fast Video Processing
 # ==========================
 def process_video(video_data):
     try:
@@ -218,120 +249,104 @@ 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}")
 
-        # Calculate optimal frame skipping
-        original_frame_skip = CONFIG["FRAME_SKIP"]
-        target_fps = min(max(fps / original_frame_skip, CONFIG["MIN_FRAME_RATE"]), CONFIG["MAX_FRAME_RATE"])
-        actual_frame_skip = max(1, int(fps / target_fps))
-        frames_to_process = total_frames // actual_frame_skip
-        logger.info(f"Processing strategy: Frame skip={actual_frame_skip}, Target FPS={target_fps:.1f}, Frames to process={frames_to_process}")
+        # Prepare for parallel processing
+        frame_batches = []
+        frame_indices_batches = []
+        current_batch = []
+        current_indices = []
+        
+        # Read all frames upfront for parallel processing
+        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)
+            
+            # Organize into batches
+            if len(current_batch) >= CONFIG["CHUNK_SIZE"]:
+                frame_batches.append(current_batch)
+                frame_indices_batches.append(current_indices)
+                current_batch = []
+                current_indices = []
+        
+        # Add remaining frames
+        if current_batch:
+            frame_batches.append(current_batch)
+            frame_indices_batches.append(current_indices)
 
+        cap.release()
+        
+        # Process frames in parallel
         workers = []
         violations = []
         helmet_violations = {}
         snapshots = []
         start_time = time.time()
-        processed_frames = 0
-        last_progress_update = 0
+        
+        # Use multiprocessing Pool
+        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))
+            
+            # Flatten results
+            all_detections = []
+            for batch_result in results:
+                all_detections.extend(batch_result)
 
-        # Process frames in batches
-        while True:
-            batch_frames = []
-            batch_indices = []
+        # Process detections and track workers
+        workers = []
+        for frame_idx, detections in sorted(all_detections, key=lambda x: x[0]):
+            current_time = frame_idx / fps
             
-            # 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
-                
-                batch_frames.append(frame)
-                batch_indices.append(frame_idx)
-                processed_frames += 1
-                
-                # Skip frames according to our strategy
-                for _ in range(actual_frame_skip - 1):
-                    if not cap.grab():
-                        break
+            # 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()
+
+            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"])
+                    if iou > max_iou and iou > CONFIG["IOU_THRESHOLD"]:
+                        max_iou = iou
+                        worker_id = worker["id"]
+                        workers[idx]["bbox"] = detection["bounding_box"]
+                        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"],
+                        "first_seen": current_time,
+                        "last_seen": current_time
+                    })
 
-            # Break if no more frames
-            if not batch_frames:
-                break
+                detection["worker_id"] = worker_id
 
-            # Run batch detection
-            results = model(batch_frames, device=device, conf=0.1, iou=CONFIG["IOU_THRESHOLD"], 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() - last_progress_update > 1.0:  # Update every second
-                    progress = (frame_idx / total_frames) * 100
-                    yield f"Processing video... {progress:.1f}% complete (Frame {frame_idx}/{total_frames})", "", "", "", ""
-                    last_progress_update = 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
-                    }
-
-                    # Worker tracking
-                    worker_id = None
-                    max_iou = 0
-                    for idx, worker in enumerate(workers):
-                        iou = calculate_iou(bbox, worker["bbox"])
-                        if iou > max_iou and iou > CONFIG["IOU_THRESHOLD"]:
-                            max_iou = iou
-                            worker_id = worker["id"]
-                            workers[idx]["bbox"] = bbox  # Update worker position
-                            workers[idx]["last_seen"] = current_time
-
-                    if worker_id is None:
-                        worker_id = len(workers) + 1
-                        workers.append({
-                            "id": worker_id,
-                            "bbox": bbox,
-                            "first_seen": current_time,
-                            "last_seen": current_time
-                        })
-
-                    detection["worker_id"] = worker_id
-
-                    # Special handling for helmet violations
-                    if label == "no_helmet":
-                        if worker_id not in helmet_violations:
-                            helmet_violations[worker_id] = []
-                        helmet_violations[worker_id].append(detection)
-                    else:
-                        violations.append(detection)
-
-                # Remove workers not seen recently
-                workers = [w for w in workers if current_time - w["last_seen"] < CONFIG["WORKER_TRACKING_DURATION"]]
+                # Special handling for helmet violations
+                if detection["violation"] == "no_helmet":
+                    if worker_id not in helmet_violations:
+                        helmet_violations[worker_id] = []
+                    helmet_violations[worker_id].append(detection)
+                else:
+                    violations.append(detection)
+
+            # Remove workers not seen recently
+            workers = [w for w in workers if current_time - w["last_seen"] < CONFIG["WORKER_TRACKING_DURATION"]]
 
         # Process helmet violations (require consistent detections)
         for worker_id, detections in helmet_violations.items():
@@ -341,6 +356,7 @@ def process_video(video_data):
                 violations.append(best_detection)
                 
                 # Capture snapshot for this violation
+                cap = cv2.VideoCapture(video_path)
                 cap.set(cv2.CAP_PROP_POS_FRAMES, best_detection["frame"])
                 ret, snapshot_frame = cap.read()
                 if ret:
@@ -354,8 +370,8 @@ def process_video(video_data):
                         "snapshot_path": snapshot_path,
                         "snapshot_base64": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}"
                     })
+                cap.release()
 
-        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.")
@@ -472,11 +488,11 @@ interface = gr.Interface(
         gr.Textbox(label="Salesforce Record ID"),
         gr.Textbox(label="Violation Details URL")
     ],
-    title="Worksite Safety Violation Analyzer",
-    description="Upload site videos to detect safety violations (No Helmet, No Harness, Unsafe Posture, Unsafe Zone, Improper Tool Use). Non-violations are ignored.",
+    title="Ultra-Fast Safety Violation Analyzer",
+    description="Upload site videos to detect safety violations in under 30 seconds (No Helmet, No Harness, Unsafe Posture, Unsafe Zone, Improper Tool Use).",
     allow_flagging="never"
 )
 
 if __name__ == "__main__":
-    logger.info("Launching Enhanced Safety Analyzer App...")
+    logger.info("Launching Ultra-Fast Safety Analyzer App...")
     interface.launch()