Update app.py

app.py

@@ -135,6 +135,36 @@ 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"
@@ -304,13 +334,12 @@ def process_video(video_data):
 
         logger.info(f"Video properties: {duration:.2f}s, {total_frames} frames, {fps:.1f} FPS, {width}x{height}")
 
-        # Track workers only for helmet violations
-        helmet_workers = {}  # {worker_id: {"first_detected": timestamp, "bbox": bbox}}
+        workers = []
         violations = []
+        helmet_violations = {}
         snapshots = []
         start_time = time.time()
         frame_skip = CONFIG["FRAME_SKIP"]
-        next_worker_id = 1
 
         # Process frames in batches
         while True:
@@ -354,8 +383,6 @@ 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)
@@ -365,99 +392,78 @@ def process_video(video_data):
                         continue
 
                     bbox = [round(x, 2) for x in box.xywh.cpu().numpy()[0]]
-                    
-                    # 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()
+                    detection = {
+                        "frame": frame_idx,
+                        "violation": label,
+                        "confidence": round(conf, 2),
+                        "bounding_box": bbox,
+                        "timestamp": current_time
+                    }
+
+                    # 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:
-                        # 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)
-                            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()
+                        violations.append(detection)
 
                 # Remove inactive workers
-                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]
+                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.")
 
+        # 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()
+
         # Generate results
         if not violations:
             yield "No violations detected in the video.", "Safety Score: 100%", "No snapshots captured.", "N/A", "N/A"
@@ -467,33 +473,22 @@ 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 (Helmet Only) |\n"
-        violation_table += "|------------------------|---------------|------------|--------------------------|\n"
+        violation_table = "| Violation              | Timestamp (s) | Confidence | Worker ID |\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")
-            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"
+            row = f"| {display_name:<22} | {v.get('timestamp', 0.0):.2f}        | {v.get('confidence', 0.0):.2f}      | {v.get('worker_id', 'N/A')} |\n"
             violation_table += row
 
-        # 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>"
+        snapshots_text = "\n".join(
+            f"- Snapshot for {CONFIG['DISPLAY_NAMES'].get(s['violation'], 'Unknown')} at frame {s['frame']}: ![]({s['snapshot_base64']})"
+            for s in snapshots
+        ) if snapshots else "No snapshots captured."
 
         yield (
             violation_table,
             f"Safety Score: {score}%",
-            snapshots_html,
+            snapshots_text,
             f"Salesforce Record ID: {report_id or 'N/A'}",
             final_pdf_url or "N/A"
         )
@@ -512,8 +507,8 @@ def gradio_interface(video_file):
         with open(video_file, "rb") as f:
             video_data = f.read()
 
-        for status, score, snapshots_html, record_id, details_url in process_video(video_data):
-            yield status, score, snapshots_html, record_id, details_url
+        for status, score, snapshots_text, record_id, details_url in process_video(video_data):
+            yield status, score, snapshots_text, 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.", "", ""
@@ -524,7 +519,7 @@ interface = gr.Interface(
     outputs=[
         gr.Markdown(label="Detected Safety Violations"),
         gr.Textbox(label="Compliance Score"),
-        gr.HTML(label="Violation Snapshots (Click to enlarge)"),
+        gr.Markdown(label="Snapshots"),
         gr.Textbox(label="Salesforce Record ID"),
         gr.Textbox(label="Violation Details URL")
     ],