Update app.py

app.py

@@ -49,7 +49,7 @@ CONFIG = {
         "domain": "login"
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo2/resolve/main/static/output/",
-    "FRAME_SKIP": 5,  # Process every 5th frame (balance speed vs. accuracy)
+    "FRAME_SKIP": 10,  # Increased to process every 10th frame (faster)
     "MAX_PROCESSING_TIME": 60,  # Max processing time (seconds)
     "CONFIDENCE_THRESHOLD": {  # Per-class thresholds
         "no_helmet": 0.4,
@@ -59,7 +59,7 @@ CONFIG = {
         "improper_tool_use": 0.35
     },
     "IOU_THRESHOLD": 0.4,  # For worker tracking
-    "MIN_VIOLATION_FRAMES": 3  # Min frames to confirm a violation
+    "MIN_VIOLATION_FRAMES": 2  # Reduced to 2 frames for faster confirmation
 }
 
 # Setup logging
@@ -239,6 +239,7 @@ def process_video(video_path):
         snapshots = []
         workers = []
         snapshot_taken = {label: False for label in CONFIG["VIOLATION_LABELS"].values()}
+        frames_to_save = []  # Store frames for snapshot saving later
         
         while cap.isOpened():
             ret, frame = cap.read()
@@ -271,37 +272,18 @@ def process_video(video_path):
                         "timestamp": current_time
                     }
                     
-                    # Track worker
-                    matched_worker = None
-                    max_iou = 0
-                    for worker in workers:
-                        iou = calculate_iou(worker["bbox"], bbox)
-                        if iou > max_iou and iou > CONFIG["IOU_THRESHOLD"]:
-                            max_iou = iou
-                            matched_worker = worker
-                    
-                    if matched_worker:
-                        worker_id = matched_worker["id"]
-                        matched_worker["bbox"] = bbox
-                    else:
-                        worker_id = len(workers) + 1
-                        workers.append({"id": worker_id, "bbox": bbox})
-                    
+                    # Simplified worker tracking
+                    worker_id = len(workers) + 1  # Assign new ID without IoU for speed
+                    workers.append({"id": worker_id, "bbox": bbox})
                     detection["worker_id"] = worker_id
                     violations.append(detection)
                     
-                    # Capture snapshot if first detection of this type
+                    # Store frame for snapshot if first detection of this type
                     if not snapshot_taken[label]:
-                        snapshot_path = os.path.join(
-                            CONFIG["OUTPUT_DIR"],
-                            f"{label}_{frame_count}.jpg"
-                        )
-                        cv2.imwrite(snapshot_path, draw_detections(frame.copy(), [detection]))
-                        snapshots.append({
-                            "violation": label,
-                            "frame": frame_count,
-                            "path": snapshot_path,
-                            "url": f"{CONFIG['PUBLIC_URL_BASE']}{os.path.basename(snapshot_path)}"
+                        frames_to_save.append({
+                            "frame": frame.copy(),
+                            "detection": detection,
+                            "label": label
                         })
                         snapshot_taken[label] = True
             
@@ -309,6 +291,23 @@ def process_video(video_path):
         
         cap.release()
         
+        # Save snapshots after processing all frames
+        for item in frames_to_save:
+            frame = item["frame"]
+            detection = item["detection"]
+            label = item["label"]
+            snapshot_path = os.path.join(
+                CONFIG["OUTPUT_DIR"],
+                f"{label}_{detection['frame']}.jpg"
+            )
+            cv2.imwrite(snapshot_path, draw_detections(frame, [detection]))
+            snapshots.append({
+                "violation": label,
+                "frame": detection["frame"],
+                "path": snapshot_path,
+                "url": f"{CONFIG['PUBLIC_URL_BASE']}{os.path.basename(snapshot_path)}"
+            })
+        
         # Filter violations (require min frames)
         filtered_violations = []
         violation_counts = {}