Update app.py

app.py

@@ -64,8 +64,7 @@ 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
-    "VIOLATION_COOLDOWN": 5.0  # Seconds before same violation can be reported again for same worker
+    "PARALLEL_WORKERS": max(1, cpu_count() - 1)  # Use all CPU cores except one
 }
 
 # Setup logging
@@ -305,12 +304,13 @@ def process_video(video_data):
 
         logger.info(f"Video properties: {duration:.2f}s, {total_frames} frames, {fps:.1f} FPS, {width}x{height}")
 
-        workers = []
+        # Track workers only for helmet violations
+        helmet_workers = {}  # {worker_id: {"first_detected": timestamp, "bbox": bbox}}
         violations = []
-        violation_history = {}  # Track when violations were last reported for each worker
         snapshots = []
         start_time = time.time()
         frame_skip = CONFIG["FRAME_SKIP"]
+        next_worker_id = 1
 
         # Process frames in batches
         while True:
@@ -354,6 +354,8 @@ def process_video(video_data):
 
                 # Process detections in this frame
                 boxes = result.boxes
+                frame_violations = set()  # Track violations in this frame to avoid duplicates
+                
                 for box in boxes:
                     cls = int(box.cls)
                     conf = float(box.conf)
@@ -364,80 +366,92 @@ def process_video(video_data):
 
                     bbox = [round(x, 2) for x in box.xywh.cpu().numpy()[0]]
                     
-                    # 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 > 0.4:  # IOU threshold
-                            max_iou = iou
-                            worker_id = worker["id"]
-                            workers[idx]["bbox"] = bbox
-                            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
-                        })
-
-                    # Check if we should report this violation (cooldown period)
-                    violation_key = f"{worker_id}_{label}"
-                    last_reported = violation_history.get(violation_key, 0)
-                    
-                    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:
+                    # For no_helmet violations, track workers and only record first occurrence
+                    if label == "no_helmet":
+                        # Check if this is a known worker
+                        worker_id = None
+                        for w_id, worker in helmet_workers.items():
+                            iou = calculate_iou(bbox, worker["bbox"])
+                            if iou > 0.4:  # IOU threshold
+                                worker_id = w_id
+                                # Update worker's position
+                                helmet_workers[w_id]["bbox"] = bbox
+                                helmet_workers[w_id]["last_seen"] = current_time
+                                break
+                        
+                        # If new worker, assign ID and record first violation
+                        if worker_id is None:
+                            worker_id = next_worker_id
+                            next_worker_id += 1
+                            helmet_workers[worker_id] = {
+                                "bbox": bbox,
+                                "first_seen": current_time,
+                                "last_seen": current_time
+                            }
+                            
+                            # Only record first violation for this worker
+                            detection = {
+                                "frame": frame_idx,
+                                "violation": label,
+                                "confidence": round(conf, 2),
+                                "bounding_box": bbox,
+                                "timestamp": current_time,
+                                "worker_id": worker_id
+                            }
+                            violations.append(detection)
+                            
+                            # Capture snapshot
+                            cap_snapshot = cv2.VideoCapture(video_path)
+                            cap_snapshot.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
+                            ret, snapshot_frame = cap_snapshot.read()
+                            if ret:
+                                snapshot_frame = draw_detections(snapshot_frame, [detection])
+                                snapshot_filename = f"no_helmet_{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": frame_idx,
+                                    "worker_id": worker_id,
+                                    "snapshot_path": snapshot_path,
+                                    "snapshot_url": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}"
+                                })
+                            cap_snapshot.release()
+                    else:
+                        # For other violations, only record if not already detected in this frame
+                        if label not in frame_violations:
+                            detection = {
+                                "frame": frame_idx,
+                                "violation": label,
+                                "confidence": round(conf, 2),
+                                "bounding_box": bbox,
+                                "timestamp": current_time
+                            }
                             violations.append(detection)
-                            violation_history[violation_key] = current_time
+                            frame_violations.add(label)
+                            
+                            # Capture snapshot for first occurrence of this violation type
+                            cap_snapshot = cv2.VideoCapture(video_path)
+                            cap_snapshot.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
+                            ret, snapshot_frame = cap_snapshot.read()
+                            if ret:
+                                snapshot_frame = draw_detections(snapshot_frame, [detection])
+                                snapshot_filename = f"{label}_{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,
+                                    "snapshot_path": snapshot_path,
+                                    "snapshot_url": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}"
+                                })
+                            cap_snapshot.release()
 
                 # Remove inactive workers
-                workers = [w for w in workers if current_time - w["last_seen"] < CONFIG["WORKER_TRACKING_DURATION"]]
+                inactive_workers = [w_id for w_id, worker in helmet_workers.items() 
+                                  if current_time - worker["last_seen"] > CONFIG["WORKER_TRACKING_DURATION"]]
+                for w_id in inactive_workers:
+                    del helmet_workers[w_id]
 
         cap.release()
         os.remove(video_path)
@@ -453,32 +467,33 @@ 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 | Snapshot |\n"
-        violation_table += "|------------------------|---------------|------------|-----------|----------|\n"
+        violation_table = "| Violation              | Timestamp (s) | Confidence | Worker ID (Helmet Only) |\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")
-            # 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"
+            worker_id = v.get("worker_id", "N/A") if v.get("violation") == "no_helmet" else "N/A"
+            row = f"| {display_name:<22} | {v.get('timestamp', 0.0):.2f}        | {v.get('confidence', 0.0):.2f}      | {worker_id} |\n"
             violation_table += row
 
-        snapshots_text = "\n".join(
-            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."
+        # Create HTML for snapshots with clickable links
+        snapshots_html = "<div style='display: flex; flex-wrap: wrap; gap: 10px;'>"
+        for s in snapshots:
+            display_name = CONFIG["DISPLAY_NAMES"].get(s["violation"], "Unknown")
+            worker_text = f"Worker {s['worker_id']}" if "worker_id" in s else ""
+            snapshots_html += f"""
+            <div style='text-align: center; margin: 10px;'>
+                <a href='{s['snapshot_url']}' target='_blank'>
+                    <img src='{s['snapshot_url']}' style='max-width: 200px; max-height: 150px;'/>
+                </a>
+                <p>{display_name} at frame {s['frame']} {worker_text}</p>
+            </div>
+            """
+        snapshots_html += "</div>"
 
         yield (
             violation_table,
             f"Safety Score: {score}%",
-            snapshots_text,
+            snapshots_html,
             f"Salesforce Record ID: {report_id or 'N/A'}",
             final_pdf_url or "N/A"
         )
@@ -497,8 +512,8 @@ def gradio_interface(video_file):
         with open(video_file, "rb") as f:
             video_data = f.read()
 
-        for status, score, snapshots_text, record_id, details_url in process_video(video_data):
-            yield status, score, snapshots_text, record_id, details_url
+        for status, score, snapshots_html, record_id, details_url in process_video(video_data):
+            yield status, score, snapshots_html, record_id, details_url
     except Exception as e:
         logger.error(f"Error in Gradio interface: {e}", exc_info=True)
         yield f"Error: {str(e)}", "", "Error in processing.", "", ""
@@ -509,7 +524,7 @@ interface = gr.Interface(
     outputs=[
         gr.Markdown(label="Detected Safety Violations"),
         gr.Textbox(label="Compliance Score"),
-        gr.Markdown(label="Snapshots"),
+        gr.HTML(label="Violation Snapshots (Click to enlarge)"),
         gr.Textbox(label="Salesforce Record ID"),
         gr.Textbox(label="Violation Details URL")
     ],