Update app.py

app.py

@@ -18,8 +18,8 @@ from retrying import retry
 # Configuration
 # ==========================
 CONFIG = {
-    "MODEL_PATH": "yolov8_safety.pt",  # Make sure this file exists in your directory
-    "FALLBACK_MODEL": "yolov8n.pt",  # Fallback model for testing
+    "MODEL_PATH": "yolov8_safety.pt",
+    "FALLBACK_MODEL": "yolov8n.pt",
     "OUTPUT_DIR": "static/output",
     "VIOLATION_LABELS": {
         0: "no_helmet",
@@ -40,7 +40,8 @@ CONFIG = {
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo2/resolve/main/static/output/",
     "FRAME_SKIP": 15,
     "MAX_PROCESSING_TIME": 30,
-    "CONFIDENCE_THRESHOLD": 0.5
+    "CONFIDENCE_THRESHOLD": 0.5,
+    "IOU_THRESHOLD": 0.5  # Added for worker tracking
 }
 
 # Setup logging
@@ -54,14 +55,12 @@ logger.info(f"Using device: {device}")
 
 def load_model():
     try:
-        # Check if the model file exists
         if os.path.isfile(CONFIG["MODEL_PATH"]):
             model_path = CONFIG["MODEL_PATH"]
             logger.info(f"Model loaded: {model_path}")
         else:
             model_path = CONFIG["FALLBACK_MODEL"]
             logger.warning("Using fallback model. Detection accuracy may be poor. Train yolov8_safety.pt for best results.")
-            # Download fallback model if necessary
             if not os.path.isfile(model_path):
                 logger.info(f"Downloading fallback model: {model_path}")
                 torch.hub.download_url_to_file('https://github.com/ultralytics/assets/releases/download/v8.3.0/yolov8n.pt', model_path)
@@ -73,6 +72,33 @@ def load_model():
 
 model = load_model()
 
+# ==========================
+# Helper Functions
+# ==========================
+def calculate_iou(box1, box2):
+    """Calculate Intersection over Union (IoU) for two bounding boxes."""
+    x1, y1, w1, h1 = box1
+    x2, y2, w2, h2 = box2
+    
+    # Convert to top-left and bottom-right coordinates
+    x1_min, y1_min = x1 - w1/2, y1 - h1/2
+    x1_max, y1_max = x1 + w1/2, y1 + h1/2
+    x2_min, y2_min = x2 - w2/2, y2 - h2/2
+    x2_max, y2_max = x2 + w2/2, y2 + h2/2
+    
+    # Calculate intersection
+    x_min = max(x1_min, x2_min)
+    y_min = max(y1_min, y2_min)
+    x_max = min(x1_max, x2_max)
+    y_max = min(y1_max, y2_max)
+    
+    intersection = max(0, x_max - x_min) * max(0, y_max - y_min)
+    area1 = w1 * h1
+    area2 = w2 * h2
+    union = area1 + area2 - intersection
+    
+    return intersection / union if union > 0 else 0
+
 # ==========================
 # Salesforce Integration
 # ==========================
@@ -81,7 +107,7 @@ def connect_to_salesforce():
     try:
         sf = Salesforce(**CONFIG["SF_CREDENTIALS"])
         logger.info("Connected to Salesforce")
-        sf.describe()  # verify connection and metadata fetch
+        sf.describe()
         return sf
     except Exception as e:
         logger.error(f"Salesforce connection failed: {e}")
@@ -229,7 +255,7 @@ def process_video(video_data):
         fps = video.get(cv2.CAP_PROP_FPS)
 
         snapshot_taken = {label: False for label in CONFIG["VIOLATION_LABELS"].values()}
-        detected_violations = {}
+        workers = []  # List to track workers: [{"id": int, "violations": set(), "bbox": list, "last_frame": int}]
 
         while True:
             ret, frame = video.read()
@@ -245,6 +271,7 @@ def process_video(video_data):
                 break
 
             results = model(frame, device=device)
+            current_detections = []
             for result in results:
                 for box in result.boxes:
                     cls, conf = int(box.cls), float(box.conf)
@@ -252,31 +279,82 @@ def process_video(video_data):
                     if label not in CONFIG["VIOLATION_LABELS"].values() or conf < CONFIG["CONFIDENCE_THRESHOLD"]:
                         continue
 
-                    violation_key = f"{label}_{frame_count}"
-                    if violation_key in detected_violations:
-                        continue
-
-                    detected_violations[violation_key] = {
-                        "frame": frame_count,
+                    bbox = [round(x, 2) for x in box.xywh.cpu().numpy()[0]]
+                    current_detections.append({
                         "violation": label,
                         "confidence": round(conf, 2),
-                        "bounding_box": [round(x, 2) for x in box.xywh.cpu().numpy()[0]],
-                        "timestamp": frame_count / fps
-                    }
-                    violations.append(detected_violations[violation_key])
-
-                    if not snapshot_taken[label]:
-                        snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], f"snapshot_{frame_count}_{label}.jpg")
+                        "bounding_box": bbox,
+                        "timestamp": frame_count / fps,
+                        "frame": frame_count
+                    })
+
+            # Assign detections to workers
+            for detection in current_detections:
+                matched_worker = None
+                max_iou = 0
+                for worker in workers:
+                    iou = calculate_iou(detection["bounding_box"], worker["bbox"])
+                    if iou > max_iou and iou > CONFIG["IOU_THRESHOLD"]:
+                        max_iou = iou
+                        matched_worker = worker
+
+                if matched_worker:
+                    # Update existing worker
+                    if detection["violation"] not in matched_worker["violations"]:
+                        matched_worker["violations"].add(detection["violation"])
+                        violations.append({
+                            "frame": frame_count,
+                            "violation": detection["violation"],
+                            "confidence": detection["confidence"],
+                            "bounding_box": detection["bounding_box"],
+                            "timestamp": detection["timestamp"],
+                            "worker_id": matched_worker["id"]
+                        })
+                        # Take snapshot if not already taken for this violation type
+                        if not snapshot_taken[detection["violation"]]:
+                            snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], f"snapshot_{frame_count}_{detection['violation']}.jpg")
+                            cv2.imwrite(snapshot_path, frame)
+                            with open(snapshot_path, "rb") as img_file:
+                                img_base64 = base64.b64encode(img_file.read()).decode('utf-8')
+                            snapshots.append({
+                                "violation": detection["violation"],
+                                "frame": frame_count,
+                                "snapshot_url": f"{CONFIG['PUBLIC_URL_BASE']}{os.path.basename(snapshot_path)}",
+                                "snapshot_base64": f"data:image/jpeg;base64,{img_base64}"
+                            })
+                            snapshot_taken[detection["violation"]] = True
+                    matched_worker["bbox"] = detection["bounding_box"]
+                    matched_worker["last_frame"] = frame_count
+                else:
+                    # New worker
+                    worker_id = len(workers) + 1
+                    workers.append({
+                        "id": worker_id,
+                        "violations": {detection["violation"]},
+                        "bbox": detection["bounding_box"],
+                        "last_frame": frame_count
+                    })
+                    violations.append({
+                        "frame": frame_count,
+                        "violation": detection["violation"],
+                        "confidence": detection["confidence"],
+                        "bounding_box": detection["bounding_box"],
+                        "timestamp": detection["timestamp"],
+                        "worker_id": worker_id
+                    })
+                    # Take snapshot if not already taken for this violation type
+                    if not snapshot_taken[detection["violation"]]:
+                        snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], f"snapshot_{frame_count}_{detection['violation']}.jpg")
                         cv2.imwrite(snapshot_path, frame)
                         with open(snapshot_path, "rb") as img_file:
                             img_base64 = base64.b64encode(img_file.read()).decode('utf-8')
                         snapshots.append({
-                            "violation": label,
+                            "violation": detection["violation"],
                             "frame": frame_count,
                             "snapshot_url": f"{CONFIG['PUBLIC_URL_BASE']}{os.path.basename(snapshot_path)}",
                             "snapshot_base64": f"data:image/jpeg;base64,{img_base64}"
                         })
-                        snapshot_taken[label] = True
+                        snapshot_taken[detection["violation"]] = True
 
             frame_count += 1
 
@@ -334,13 +412,13 @@ def gradio_interface(video_file):
 
         violation_table = "No violations detected."
         if result["violations"]:
-            header = "| Violation              | Timestamp (s) | Confidence | Bounding Box             |\n"
-            separator = "|------------------------|---------------|------------|--------------------------|\n"
+            header = "| Violation              | Timestamp (s) | Confidence | Worker ID |\n"
+            separator = "|------------------------|---------------|------------|-----------|\n"
             rows = []
             violation_name_map = CONFIG["DISPLAY_NAMES"]
             for v in result["violations"]:
                 display_name = violation_name_map.get(v["violation"], v["violation"])
-                row = f"| {display_name:<22} | {v['timestamp']:.2f}        | {v['confidence']:.2f}      | {v['bounding_box']} |"
+                row = f"| {display_name:<22} | {v['timestamp']:.2f}        | {v['confidence']:.2f}      | {v['worker_id']} |"
                 rows.append(row)
             violation_table = header + separator + "\n".join(rows)
 
@@ -380,4 +458,4 @@ interface = gr.Interface(
 
 if __name__ == "__main__":
     logger.info("Launching Safety Analyzer App...")
-    interface.launch()
+    interface.launch()