Update app.py

app.py

@@ -64,7 +64,8 @@ CONFIG = {
     "MAX_PROCESSING_TIME": 60,  # 1 minute limit
     "FRAME_SKIP": 2,  # Process every 2nd frame for speed
     "BATCH_SIZE": 16,  # Frames per batch
-    "PARALLEL_WORKERS": max(1, cpu_count() - 1)  # Use all CPU cores except one
+    "PARALLEL_WORKERS": max(1, cpu_count() - 1),  # Use all CPU cores except one
+    "VIOLATION_COOLDOWN": 5.0  # Seconds before same violation can be detected again for same worker
 }
 
 # Setup logging
@@ -135,36 +136,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"
@@ -336,10 +307,12 @@ def process_video(video_data):
 
         workers = []
         violations = []
-        helmet_violations = {}
         snapshots = []
         start_time = time.time()
         frame_skip = CONFIG["FRAME_SKIP"]
+        
+        # Track active violations to prevent duplicate detections
+        active_violations = {}  # Format: {worker_id: {violation_type: last_detection_time}}
 
         # Process frames in batches
         while True:
@@ -392,14 +365,7 @@ def process_video(video_data):
                         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
@@ -420,49 +386,59 @@ def process_video(video_data):
                             "last_seen": current_time
                         })
 
-                    detection["worker_id"] = worker_id
+                    # Check if this worker already has this violation type in cooldown period
+                    if worker_id in active_violations:
+                        if label in active_violations[worker_id]:
+                            last_detection_time = active_violations[worker_id][label]
+                            if current_time - last_detection_time < CONFIG["VIOLATION_COOLDOWN"]:
+                                continue  # Skip this detection as it's within cooldown period
 
-                    # Track helmet violations with stricter criteria
-                    if detection["violation"] == "no_helmet":
-                        # Only include high-confidence no_helmet detections
-                        if conf >= CONFIG["CONFIDENCE_THRESHOLDS"]["no_helmet"]:
-                            if worker_id not in helmet_violations:
-                                helmet_violations[worker_id] = []
-                            helmet_violations[worker_id].append(detection)
-                    else:
-                        violations.append(detection)
-
-                # 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.")
+                    # Create detection record
+                    detection = {
+                        "frame": frame_idx,
+                        "violation": label,
+                        "confidence": round(conf, 2),
+                        "bounding_box": bbox,
+                        "timestamp": current_time,
+                        "worker_id": worker_id
+                    }
 
-        # Confirm helmet violations (require multiple detections)
-        for worker_id, detections in helmet_violations.items():
-            if len(detections) >= CONFIG["MIN_VIOLATION_FRAMES"]:
-                # Select the detection with the highest confidence
-                best_detection = max(detections, key=lambda x: x["confidence"])
-                violations.append(best_detection)
-
-                # Capture snapshot for confirmed no_helmet 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"
+                    # Add to violations list
+                    violations.append(detection)
+                    
+                    # Update active violations
+                    if worker_id not in active_violations:
+                        active_violations[worker_id] = {}
+                    active_violations[worker_id][label] = current_time
+
+                    # Capture snapshot for this violation
+                    snapshot_frame = batch_frames[i].copy()
+                    snapshot_frame = draw_detections(snapshot_frame, [detection])
+                    snapshot_filename = f"{label}_{worker_id}_{frame_idx}.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"],
+                        "violation": label,
+                        "frame": frame_idx,
+                        "worker_id": worker_id,
                         "snapshot_path": snapshot_path,
                         "snapshot_base64": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}"
                     })
-                cap.release()
+
+                # Remove inactive workers and their violations
+                workers = [w for w in workers if current_time - w["last_seen"] < CONFIG["WORKER_TRACKING_DURATION"]]
+                
+                # Clean up active_violations for workers no longer tracked
+                active_violations = {
+                    worker_id: violations 
+                    for worker_id, violations in active_violations.items() 
+                    if any(w["id"] == worker_id for w in workers)
+                }
+
+        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:
@@ -481,7 +457,7 @@ def process_video(video_data):
             violation_table += row
 
         snapshots_text = "\n".join(
-            f"- Snapshot for {CONFIG['DISPLAY_NAMES'].get(s['violation'], 'Unknown')} at frame {s['frame']}: ![]({s['snapshot_base64']})"
+            f"- Snapshot for {CONFIG['DISPLAY_NAMES'].get(s['violation'], 'Unknown')} (Worker {s['worker_id']}) at frame {s['frame']}: ![]({s['snapshot_base64']})"
             for s in snapshots
         ) if snapshots else "No snapshots captured."