Update app.py

app.py

@@ -4,6 +4,7 @@ import gradio as gr
 import torch
 import numpy as np
 from ultralytics import YOLO
+from bytetrack import BYTETracker  # Added ByteTrack
 import time
 from simple_salesforce import Salesforce
 from reportlab.lib.pagesizes import letter
@@ -60,11 +61,14 @@ CONFIG = {
         "improper_tool_use": 0.4
     },
     "MIN_VIOLATION_FRAMES": 3,
-    "WORKER_TRACKING_DURATION": 3.0,
+    "WORKER_TRACKING_DURATION": 5.0,  # Increased to 5s for better continuity
     "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
+    "FRAME_SKIP": 1,  # Process all frames to avoid missing violations
+    "BATCH_SIZE": 32,  # Increased for performance
+    "PARALLEL_WORKERS": max(1, cpu_count() - 1),  # Use all CPU cores except one
+    "TRACK_BUFFER": 30,  # Frames to keep a track alive
+    "TRACK_THRESH": 0.4,  # Tracking confidence threshold
+    "MATCH_THRESH": 0.8  # IOU threshold for matching
 }
 
 # Setup logging
@@ -116,54 +120,17 @@ def draw_detections(frame, detections):
         cv2.putText(frame, display_text, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
     return frame
 
-def calculate_iou(box1, box2):
-    x1, y1, w1, h1 = box1
-    x2, y2, w2, h2 = box2
-    
-    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)
-    y_bottom = min(y1 + h1/2, y2 + h2/2)
-    
-    if x_right < x_left or y_bottom < y_top:
-        return 0.0
-    
-    intersection_area = (x_right - x_left) * (y_bottom - y_top)
-    box1_area = w1 * h1
-    box2_area = w2 * h2
-    union_area = box1_area + box2_area - intersection_area
-    
-    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 calculate_safety_score(violations):
+    penalties = {
+        "no_helmet": 25,
+        "no_harness": 30,
+        "unsafe_posture": 20,
+        "unsafe_zone": 35,
+        "improper_tool_use": 25
+    }
+    total_penalty = sum(penalties.get(v.get("violation", "Unknown"), 0) for v in violations)
+    score = 100 - total_penalty
+    return max(score, 0)
 
 def generate_violation_pdf(violations, score):
     try:
@@ -193,7 +160,7 @@ def generate_violation_pdf(violations, score):
         else:
             for v in violations:
                 display_name = CONFIG["DISPLAY_NAMES"].get(v.get("violation", "Unknown"), "Unknown")
-                text = f"{display_name} at {v.get('timestamp', 0.0):.2f}s (Confidence: {v.get('confidence', 0.0):.2f})"
+                text = f"{display_name} from {v.get('start_timestamp', 0.0):.2f}s to {v.get('end_timestamp', 0.0):.2f}s (Confidence: {v.get('confidence', 0.0):.2f}, Worker ID: {v.get('worker_id', 'N/A')})"
                 c.drawString(1 * inch, y_position, text)
                 y_position -= 0.3 * inch
                 if y_position < 1 * inch:
@@ -214,18 +181,6 @@ def generate_violation_pdf(violations, score):
         logger.error(f"Error generating PDF: {e}")
         return "", "", None
 
-def calculate_safety_score(violations):
-    penalties = {
-        "no_helmet": 25,
-        "no_harness": 30,
-        "unsafe_posture": 20,
-        "unsafe_zone": 35,
-        "improper_tool_use": 25
-    }
-    total_penalty = sum(penalties.get(v.get("violation", "Unknown"), 0) for v in violations)
-    score = 100 - total_penalty
-    return max(score, 0)
-
 # ==========================
 # Salesforce Integration
 # ==========================
@@ -268,7 +223,7 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
     try:
         sf = connect_to_salesforce()
         violations_text = "\n".join(
-            f"{CONFIG['DISPLAY_NAMES'].get(v.get('violation', 'Unknown'), 'Unknown')} at {v.get('timestamp', 0.0):.2f}s (Confidence: {v.get('confidence', 0.0):.2f})"
+            f"{CONFIG['DISPLAY_NAMES'].get(v.get('violation', 'Unknown'), 'Unknown')} from {v.get('start_timestamp', 0.0):.2f}s to {v.get('end_timestamp', 0.0):.2f}s (Confidence: {v.get('confidence', 0.0):.2f}, Worker ID: {v.get('worker_id', 'N/A')})"
             for v in violations
         ) or "No violations detected."
         pdf_url = f"{CONFIG['PUBLIC_URL_BASE']}{os.path.basename(pdf_path)}" if pdf_path else ""
@@ -325,18 +280,22 @@ def process_video(video_data):
 
         # Get video properties
         total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-        fps = cap.get(cv2.CAP_PROP_FPS)
-        if fps <= 0:
-            fps = 30
+        fps = cap.get(cv2.CAP_PROP_FPS) or 30
         duration = total_frames / fps
         width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
         height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-
         logger.info(f"Video properties: {duration:.2f}s, {total_frames} frames, {fps:.1f} FPS, {width}x{height}")
 
-        workers = []
-        violations = []
-        helmet_violations = {}
+        # Initialize ByteTrack
+        tracker = BYTETracker(
+            track_thresh=CONFIG["TRACK_THRESH"],
+            track_buffer=CONFIG["TRACK_BUFFER"],
+            match_thresh=CONFIG["MATCH_THRESH"],
+            frame_rate=fps
+        )
+
+        # Track violations by worker ID and type
+        violation_tracker = {}  # {worker_id: {violation_type: [detections]}}
         snapshots = []
         start_time = time.time()
         frame_skip = CONFIG["FRAME_SKIP"]
@@ -375,14 +334,15 @@ def process_video(video_data):
             for i, (result, frame_idx) in enumerate(zip(results, batch_indices)):
                 current_time = frame_idx / fps
                 
-                # Update progress periodically
-                if time.time() - start_time > 1.0:  # Update every second
+                # Update progress
+                if time.time() - start_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()
 
-                # Process detections in this frame
+                # Prepare detections for ByteTrack
                 boxes = result.boxes
+                track_inputs = []
                 for box in boxes:
                     cls = int(box.cls)
                     conf = float(box.conf)
@@ -391,78 +351,75 @@ def process_video(video_data):
                     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]]
+                    bbox = box.xywh.cpu().numpy()[0]
+                    track_inputs.append({
+                        "bbox": bbox,  # [x, y, w, h]
+                        "conf": conf,
+                        "cls": cls
+                    })
+
+                # Update tracker
+                tracked_objects = tracker.update(
+                    np.array([t["bbox"] for t in track_inputs]),
+                    np.array([t["conf"] for t in track_inputs]),
+                    np.array([t["cls"] for t in track_inputs])
+                )
+
+                # Process tracked objects
+                for obj, track_input in zip(tracked_objects, track_inputs):
+                    worker_id = obj.id
+                    label = CONFIG["VIOLATION_LABELS"].get(int(obj.cls), None)
+                    bbox = track_input["bbox"]
+                    conf = track_input["conf"]
+                    
                     detection = {
                         "frame": frame_idx,
                         "violation": label,
                         "confidence": round(conf, 2),
-                        "bounding_box": bbox,
-                        "timestamp": current_time
+                        "bounding_box": [round(x, 2) for x in bbox],
+                        "timestamp": current_time,
+                        "worker_id": worker_id
                     }
 
-                    # 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
-                        })
-
-                    detection["worker_id"] = worker_id
-
-                    # 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"]]
+                    # Track violations by worker_id and type
+                    if worker_id not in violation_tracker:
+                        violation_tracker[worker_id] = {}
+                    if label not in violation_tracker[worker_id]:
+                        violation_tracker[worker_id][label] = []
+                    violation_tracker[worker_id][label].append(detection)
 
         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.")
-
-        # 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"
-                    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()
+        logger.info(f"Processing complete in {processing_time:.2f}s")
+
+        # Consolidate violations
+        violations = []
+        for worker_id, worker_violations in violation_tracker.items():
+            for label, detections in worker_violations.items():
+                if len(detections) >= CONFIG["MIN_VIOLATION_FRAMES"]:
+                    # Select highest-confidence detection
+                    best_detection = max(detections, key=lambda x: x["confidence"])
+                    best_detection["start_timestamp"] = min(d["timestamp"] for d in detections)
+                    best_detection["end_timestamp"] = max(d["timestamp"] for d in detections)
+                    violations.append(best_detection)
+
+                    # Capture snapshot for confirmed 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"{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"],
+                            "snapshot_path": snapshot_path,
+                            "snapshot_base64": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}"
+                        })
+                    cap.release()
 
         # Generate results
         if not violations:
@@ -473,11 +430,11 @@ 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"
-        for v in sorted(violations, key=lambda x: x["timestamp"]):
+        violation_table = "| Violation              | Time Range (s) | Confidence | Worker ID |\n"
+        violation_table += "|------------------------|----------------|------------|-----------|\n"
+        for v in sorted(violations, key=lambda x: x["start_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"
+            row = f"| {display_name:<22} | {v.get('start_timestamp', 0.0):.2f}-{v.get('end_timestamp', 0.0):.2f} | {v.get('confidence', 0.0):.2f} | {v.get('worker_id', 'N/A')} |\n"
             violation_table += row
 
         snapshots_text = "\n".join(