Update app.py

app.py

@@ -65,7 +65,7 @@ CONFIG = {
     "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
-    "VIOLATION_COOLDOWN": 5.0  # Seconds before same violation can be detected again for same worker
+    "VIOLATION_COOLDOWN": 5.0  # Seconds before same violation can be reported again for same worker
 }
 
 # Setup logging
@@ -307,12 +307,10 @@ def process_video(video_data):
 
         workers = []
         violations = []
+        violation_history = {}  # Track when violations were last reported for each worker
         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:
@@ -386,54 +384,60 @@ def process_video(video_data):
                             "last_seen": current_time
                         })
 
-                    # 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
-
-                    # Create detection record
-                    detection = {
-                        "frame": frame_idx,
-                        "violation": label,
-                        "confidence": round(conf, 2),
-                        "bounding_box": bbox,
-                        "timestamp": current_time,
-                        "worker_id": worker_id
-                    }
-
-                    # Add to violations list
-                    violations.append(detection)
+                    # Check if we should report this violation (cooldown period)
+                    violation_key = f"{worker_id}_{label}"
+                    last_reported = violation_history.get(violation_key, 0)
                     
-                    # 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": label,
-                        "frame": frame_idx,
-                        "worker_id": worker_id,
-                        "snapshot_path": snapshot_path,
-                        "snapshot_base64": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}"
-                    })
-
-                # Remove inactive workers and their violations
+                    if current_time - last_reported >= CONFIG["VIOLATION_COOLDOWN"]:
+                        detection = {
+                            "frame": frame_idx,
+                            "violation": label,
+                            "confidence": round(conf, 2),
+                            "bounding_box": bbox,
+                            "timestamp": current_time,
+                            "worker_id": worker_id
+                        }
+
+                        # Track helmet violations with stricter criteria
+                        if detection["violation"] == "no_helmet":
+                            if conf >= CONFIG["CONFIDENCE_THRESHOLDS"]["no_helmet"]:
+                                # Only report if we have multiple detections
+                                if worker_id not in violation_history:
+                                    violation_history[worker_id] = []
+                                violation_history[worker_id].append(detection)
+                                
+                                # Check if we have enough detections to confirm
+                                if len(violation_history[worker_id]) >= CONFIG["MIN_VIOLATION_FRAMES"]:
+                                    # Select the detection with the highest confidence
+                                    best_detection = max(violation_history[worker_id], key=lambda x: x["confidence"])
+                                    violations.append(best_detection)
+                                    violation_history[violation_key] = current_time
+                                    
+                                    # Capture snapshot for confirmed no_helmet violation
+                                    cap_snapshot = cv2.VideoCapture(video_path)
+                                    cap_snapshot.set(cv2.CAP_PROP_POS_FRAMES, best_detection["frame"])
+                                    ret, snapshot_frame = cap_snapshot.read()
+                                    if ret:
+                                        snapshot_frame = draw_detections(snapshot_frame, [best_detection])
+                                        snapshot_filename = f"violation_{worker_id}_{label}_{best_detection['frame']}.jpg"
+                                        snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], snapshot_filename)
+                                        cv2.imwrite(snapshot_path, snapshot_frame)
+                                        snapshots.append({
+                                            "violation": label,
+                                            "frame": best_detection["frame"],
+                                            "timestamp": best_detection["timestamp"],
+                                            "worker_id": worker_id,
+                                            "snapshot_path": snapshot_path,
+                                            "snapshot_url": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}"
+                                        })
+                                    cap_snapshot.release()
+                                    del violation_history[worker_id]  # Reset for this worker
+                        else:
+                            violations.append(detection)
+                            violation_history[violation_key] = current_time
+
+                # Remove inactive workers
                 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)
@@ -449,16 +453,26 @@ def process_video(video_data):
         pdf_path, pdf_url, pdf_file = generate_violation_pdf(violations, score)
         report_id, final_pdf_url = push_report_to_salesforce(violations, score, pdf_path, pdf_file)
 
-        violation_table = "| Violation              | Timestamp (s) | Confidence | Worker ID |\n"
-        violation_table += "|------------------------|---------------|------------|-----------|\n"
+        violation_table = "| Violation              | Timestamp (s) | Confidence | Worker ID | Snapshot |\n"
+        violation_table += "|------------------------|---------------|------------|-----------|----------|\n"
         for v in sorted(violations, key=lambda x: x["timestamp"]):
             display_name = CONFIG["DISPLAY_NAMES"].get(v.get("violation", "Unknown"), "Unknown")
-            row = f"| {display_name:<22} | {v.get('timestamp', 0.0):.2f}        | {v.get('confidence', 0.0):.2f}      | {v.get('worker_id', 'N/A')} |\n"
+            # Find matching snapshot if exists
+            snapshot_markdown = ""
+            for s in snapshots:
+                if (s["worker_id"] == v.get("worker_id") and 
+                    s["violation"] == v.get("violation") and 
+                    abs(s["timestamp"] - v.get("timestamp", 0)) < 0.5):
+                    snapshot_markdown = f"[View]({s['snapshot_url']})"
+                    break
+            
+            row = f"| {display_name:<22} | {v.get('timestamp', 0.0):.2f}        | {v.get('confidence', 0.0):.2f}      | {v.get('worker_id', 'N/A')} | {snapshot_markdown} |\n"
             violation_table += row
 
         snapshots_text = "\n".join(
-            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
+            f"### {CONFIG['DISPLAY_NAMES'].get(s['violation'], 'Unknown')} (Worker {s['worker_id']}) at {s['timestamp']:.2f}s\n"
+            f"![Violation]({s['snapshot_url']})"
+            for s in sorted(snapshots, key=lambda x: x["timestamp"])
         ) if snapshots else "No snapshots captured."
 
         yield (