Update app.py

app.py

@@ -53,19 +53,19 @@ CONFIG = {
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo2/resolve/main/static/output/",
     "CONFIDENCE_THRESHOLDS": {
-        "no_helmet": 0.6,
-        "no_harness": 0.15,
-        "unsafe_posture": 0.15,
-        "unsafe_zone": 0.15,
-        "improper_tool_use": 0.15
+        "no_helmet": 0.55,
+        "no_harness": 0.1,
+        "unsafe_posture": 0.1,
+        "unsafe_zone": 0.1,
+        "improper_tool_use": 0.1
     },
-    "IOU_THRESHOLD": 0.4,
-    "MIN_VIOLATION_FRAMES": 3,
-    "HELMET_CONFIDENCE_THRESHOLD": 0.65,
-    "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
+    "IOU_THRESHOLD": 0.45,
+    "MIN_VIOLATION_FRAMES": 2,
+    "HELMET_CONFIDENCE_THRESHOLD": 0.6,
+    "WORKER_TRACKING_DURATION": 2.5,
+    "MAX_PROCESSING_TIME": 30,
+    "PARALLEL_WORKERS": max(1, cpu_count() - 1),
+    "CHUNK_SIZE": 15  # Increased chunk size for faster processing
 }
 
 # Setup logging
@@ -122,7 +122,6 @@ def calculate_iou(box1, box2):
     x1, y1, w1, h1 = box1
     x2, y2, w2, h2 = box2
     
-    # 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)
@@ -140,7 +139,7 @@ def calculate_iou(box1, box2):
 
 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)
+    results = model(frame_batch, device=device, conf=0.05, iou=CONFIG["IOU_THRESHOLD"], verbose=False)
     
     for idx, (result, frame_idx) in enumerate(zip(results, frame_indices)):
         current_time = frame_idx / fps
@@ -152,7 +151,7 @@ def process_frame_batch(frame_batch, frame_indices, fps):
             conf = float(box.conf)
             label = CONFIG["VIOLATION_LABELS"].get(cls, None)
             
-            if label is None or conf < CONFIG["CONFIDENCE_THRESHOLDS"].get(label, 0.25):
+            if label is None or conf < CONFIG["CONFIDENCE_THRESHOLDS"].get(label, 0.2):
                 continue
 
             bbox = [round(x, 2) for x in box.xywh.cpu().numpy()[0]]
@@ -234,6 +233,7 @@ def calculate_safety_score(violations):
 # ==========================
 def process_video(video_data):
     try:
+        start_time = time.time()
         # Create temp video file
         video_path = os.path.join(CONFIG["OUTPUT_DIR"], f"temp_{int(time.time())}.mp4")
         with open(video_path, "wb") as f:
@@ -256,66 +256,74 @@ def process_video(video_data):
 
         logger.info(f"Video properties: {duration:.2f}s, {total_frames} frames, {fps:.1f} FPS, {width}x{height}")
 
-        # Prepare for parallel processing
+        # Check if processing will exceed time limit
+        if duration > CONFIG["MAX_PROCESSING_TIME"]:
+            logger.warning(f"Video duration {duration:.2f}s exceeds max processing time {CONFIG['MAX_PROCESSING_TIME']}s")
+            cap.release()
+            os.remove(video_path)
+            yield "Video duration too long. Please upload a shorter video.", "", "", "", ""
+            return
+
+        # Read all frames upfront
         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):
+        frame_count = 0
+
+        while True:
             ret, frame = cap.read()
             if not ret:
                 break
-            all_frames.append(frame)
-            all_indices.append(frame_idx)
-            
-            # Organize into batches
+            current_batch.append(frame)
+            current_indices.append(frame_count)
+            frame_count += 1
+
             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()
-        
-        # Use multiprocessing Pool
+        last_progress_time = start_time
+
         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
+            # Flatten and sort results
             all_detections = []
             for batch_result in results:
                 all_detections.extend(batch_result)
+            all_detections.sort(key=lambda x: x[0])
 
-        # Process detections and track workers
+        # Worker tracking
         workers = []
-        for frame_idx, detections in sorted(all_detections, key=lambda x: x[0]):
+        for frame_idx, detections in all_detections:
             current_time = frame_idx / fps
-            
-            # Update progress periodically
-            if time.time() - start_time > 1.0:  # Update every second
+
+            # Update progress every second
+            if time.time() - last_progress_time > 1.0:
                 progress = (frame_idx / total_frames) * 100
                 yield f"Processing video... {progress:.1f}% complete (Frame {frame_idx}/{total_frames})", "", "", "", ""
-                start_time = time.time()
+                last_progress_time = time.time()
+
+            # Early termination if time limit approached
+            if time.time() - start_time > CONFIG["MAX_PROCESSING_TIME"] - 2:
+                logger.warning("Approaching max processing time, terminating early")
+                break
 
             for detection in detections:
-                # Worker tracking
                 worker_id = None
                 max_iou = 0
                 for idx, worker in enumerate(workers):
@@ -337,7 +345,6 @@ def process_video(video_data):
 
                 detection["worker_id"] = worker_id
 
-                # Special handling for helmet violations
                 if detection["violation"] == "no_helmet":
                     if worker_id not in helmet_violations:
                         helmet_violations[worker_id] = []
@@ -345,38 +352,36 @@ def process_video(video_data):
                 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)
+        # Process helmet violations
         for worker_id, detections in helmet_violations.items():
             if len(detections) >= CONFIG["MIN_VIOLATION_FRAMES"]:
-                # Find the detection with highest confidence
                 best_detection = max(detections, key=lambda x: x["confidence"])
-                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:
-                    snapshot_frame = draw_detections(snapshot_frame, [best_detection])
-                    snapshot_filename = f"no_helmet_{best_detection['frame']}.jpg"
-                    snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], snapshot_filename)
-                    cv2.imwrite(snapshot_path, snapshot_frame)
-                    snapshots.append({
-                        "violation": "no_helmet",
-                        "frame": best_detection["frame"],
-                        "snapshot_path": snapshot_path,
-                        "snapshot_base64": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}"
-                    })
-                cap.release()
+                if best_detection["confidence"] >= CONFIG["HELMET_CONFIDENCE_THRESHOLD"]:
+                    violations.append(best_detection)
+                    
+                    # Capture snapshot
+                    cap = cv2.VideoCapture(video_path)
+                    cap.set(cv2.CAP_PROP_POS_FRAMES, best_detection["frame"])
+                    ret, snapshot_frame = cap.read()
+                    if ret:
+                        snapshot_frame = draw_detections(snapshot_frame, [best_detection])
+                        snapshot_filename = f"no_helmet_{best_detection['frame']}.jpg"
+                        snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], snapshot_filename)
+                        cv2.imwrite(snapshot_path, snapshot_frame)
+                        snapshots.append({
+                            "violation": "no_helmet",
+                            "frame": best_detection["frame"],
+                            "snapshot_path": snapshot_path,
+                            "snapshot_base64": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}"
+                        })
+                    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"
             return
@@ -384,7 +389,6 @@ def process_video(video_data):
         score = calculate_safety_score(violations)
         pdf_path, pdf_url, pdf_file = generate_violation_pdf(violations, score)
 
-        # Generate violation table
         violation_table = "| Violation              | Timestamp (s) | Confidence | Worker ID |\n"
         violation_table += "|------------------------|---------------|------------|-----------|\n"
         for v in sorted(violations, key=lambda x: x["timestamp"]):
@@ -392,13 +396,11 @@ def process_video(video_data):
             row = f"| {display_name:<22} | {v.get('timestamp', 0.0):.2f}        | {v.get('confidence', 0.0):.2f}      | {v.get('worker_id', 'N/A')} |\n"
             violation_table += row
 
-        # Generate snapshots text
         snapshots_text = "\n".join(
             f"- Snapshot for {CONFIG['DISPLAY_NAMES'].get(s['violation'], 'Unknown')} at frame {s['frame']}: ![]({s['snapshot_base64']})"
             for s in snapshots
         ) if snapshots else "No snapshots captured."
 
-        # Push to Salesforce
         try:
             sf = connect_to_salesforce()
             record_data = {