Update app.py

app.py

@@ -15,7 +15,7 @@ import logging
 from retrying import retry
 
 # ==========================
-# Configuration
+# Enhanced Configuration
 # ==========================
 CONFIG = {
     "MODEL_PATH": "yolov8_safety.pt",
@@ -24,12 +24,16 @@ CONFIG = {
     "VIOLATION_LABELS": {
         0: "no_helmet",
         1: "no_harness",
-        2: "unsafe_posture"
+        2: "unsafe_posture",
+        3: "unsafe_zone",
+        4: "improper_tool_use"
     },
     "DISPLAY_NAMES": {
         "no_helmet": "No Helmet Violation",
         "no_harness": "No Harness Violation",
-        "unsafe_posture": "Unsafe Posture Violation"
+        "unsafe_posture": "Unsafe Posture Violation",
+        "unsafe_zone": "Unsafe Zone Entry",
+        "improper_tool_use": "Improper Tool Use"
     },
     "SF_CREDENTIALS": {
         "username": "prashanth1ai@safety.com",
@@ -38,10 +42,17 @@ CONFIG = {
         "domain": "login"
     },
     "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.1,  # Lowered for debugging
-    "IOU_THRESHOLD": 0.5  # For worker tracking
+    "FRAME_SKIP": 10,  # Reduced for better detection
+    "MAX_PROCESSING_TIME": 45,
+    "CONFIDENCE_THRESHOLD": {
+        "no_helmet": 0.4,
+        "no_harness": 0.35,
+        "unsafe_posture": 0.3,
+        "unsafe_zone": 0.3,
+        "improper_tool_use": 0.35
+    },
+    "IOU_THRESHOLD": 0.4,
+    "MIN_VIOLATION_DURATION": 2  # seconds
 }
 
 # Setup logging
@@ -73,7 +84,7 @@ def load_model():
 model = load_model()
 
 # ==========================
-# Helper Functions
+# Enhanced Helper Functions
 # ==========================
 def calculate_iou(box1, box2):
     """Calculate Intersection over Union (IoU) for two bounding boxes."""
@@ -99,8 +110,17 @@ def calculate_iou(box1, box2):
     
     return intersection / union if union > 0 else 0
 
+def is_violation_persistent(violation_type, worker_id, violations, fps):
+    """Check if a violation persists for the required duration."""
+    violation_times = [v['timestamp'] for v in violations 
+                      if v['violation'] == violation_type and v['worker_id'] == worker_id]
+    if len(violation_times) < 2:
+        return False
+    duration = max(violation_times) - min(violation_times)
+    return duration >= CONFIG["MIN_VIOLATION_DURATION"]
+
 # ==========================
-# Salesforce Integration
+# Salesforce Integration (unchanged)
 # ==========================
 @retry(stop_max_attempt_number=3, wait_fixed=2000)
 def connect_to_salesforce():
@@ -141,7 +161,7 @@ def generate_violation_pdf(violations, score):
         else:
             for v in violations:
                 display_name = CONFIG["DISPLAY_NAMES"].get(v["violation"], v["violation"])
-                text = f"{display_name} at {v['timestamp']:.2f}s (Confidence: {v['confidence']:.2f}, Worker ID: {v['worker_id']})"
+                text = f"{display_name} at {v['timestamp']:.2f}s (Confidence: {v['confidence']})"
                 c.drawString(1 * inch, y_position, text)
                 y_position -= 0.3 * inch
                 if y_position < 1 * inch:
@@ -190,7 +210,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['violation'], v['violation'])} at {v['timestamp']:.2f}s (Confidence: {v['confidence']:.2f}, Worker ID: {v['worker_id']})"
+            f"{CONFIG['DISPLAY_NAMES'].get(v['violation'], v['violation'])} at {v['timestamp']:.2f}s (Confidence: {v['confidence']})"
             for v in violations
         ) or "No violations detected."
         pdf_url = f"{CONFIG['PUBLIC_URL_BASE']}{os.path.basename(pdf_path)}" if pdf_path else ""
@@ -233,14 +253,23 @@ def calculate_safety_score(violations):
     penalties = {
         "no_helmet": 25,
         "no_harness": 30,
-        "unsafe_posture": 20
+        "unsafe_posture": 20,
+        "unsafe_zone": 35,
+        "improper_tool_use": 25
     }
-    total_penalty = sum(penalties.get(v["violation"], 0) for v in violations)
-    logger.info(f"Total Penalty: {total_penalty}")
+    # Count unique violations per worker
+    unique_violations = set()
+    for v in violations:
+        key = (v["worker_id"], v["violation"])
+        unique_violations.add(key)
+    
+    total_penalty = sum(penalties.get(violation, 0) for _, violation in unique_violations)
     score = 100 - total_penalty
-    logger.info(f"Calculated Score: {score}")
     return max(score, 0)
 
+# ==========================
+# Enhanced Video Processing
+# ==========================
 def process_video(video_data):
     try:
         video_path = os.path.join(CONFIG["OUTPUT_DIR"], f"temp_{int(time.time())}.mp4")
@@ -252,16 +281,19 @@ def process_video(video_data):
         if not video.isOpened():
             raise ValueError("Could not open video file")
 
-        violations, snapshots, raw_detections = [], [], []
+        violations, snapshots = [], []
         frame_count = 0
         start_time = time.time()
         fps = video.get(cv2.CAP_PROP_FPS)
+        if fps <= 0:
+            fps = 30  # Default assumption if FPS cannot be determined
 
         snapshot_taken = {label: False for label in CONFIG["VIOLATION_LABELS"].values()}
         workers = []  # List to track workers
+        violation_history = []  # Track all potential violations before filtering
 
+        logger.info(f"Processing video with FPS: {fps}")
         logger.info(f"Looking for violations: {CONFIG['VIOLATION_LABELS']}")
-        logger.info(f"Using confidence threshold: {CONFIG['CONFIDENCE_THRESHOLD']}")
 
         while True:
             ret, frame = video.read()
@@ -276,154 +308,134 @@ def process_video(video_data):
                 logger.info("Processing time limit reached")
                 break
 
+            # Run detection on this frame
             results = model(frame, device=device)
-            current_detections = []
+            current_time = frame_count / fps
             
             for result in results:
                 boxes = result.boxes
-                logger.info(f"Frame {frame_count}: Found {len(boxes)} potential detections")
+                logger.debug(f"Frame {frame_count}: Found {len(boxes)} potential detections")
                 
                 for box in boxes:
-                    cls, conf = int(box.cls), float(box.conf)
+                    cls = int(box.cls)
+                    conf = float(box.conf)
                     label = CONFIG["VIOLATION_LABELS"].get(cls, None)
                     
-                    # Log all raw detections
-                    logger.info(f"Raw Detection: class={cls}, conf={conf:.2f}, label={label}")
-                    raw_detections.append({
-                        "frame": frame_count,
-                        "class": cls,
-                        "confidence": round(conf, 2),
-                        "label": label if label in CONFIG["VIOLATION_LABELS"].values() else "unknown",
-                        "timestamp": frame_count / fps
-                    })
-                    
-                    if label not in CONFIG["VIOLATION_LABELS"].values():
-                        logger.info(f"Skipping unknown class: {cls}")
-                        continue
+                    if label is None:
+                        continue  # Skip unknown classes
                         
-                    if conf < CONFIG["CONFIDENCE_THRESHOLD"]:
-                        logger.info(f"Skipping low confidence: {conf:.2f} < {CONFIG['CONFIDENCE_THRESHOLD']}")
+                    conf_threshold = CONFIG["CONFIDENCE_THRESHOLD"].get(label, 0.3)
+                    if conf < conf_threshold:
+                        logger.debug(f"Skipping {label} with low confidence: {conf:.2f} < {conf_threshold}")
                         continue
 
                     bbox = [round(x, 2) for x in box.xywh.cpu().numpy()[0]]
-                    logger.info(f"Valid detection: {label} with confidence: {conf:.2f}")
                     
-                    current_detections.append({
+                    # Store potential violation (will filter later)
+                    violation_history.append({
+                        "frame": frame_count,
                         "violation": label,
                         "confidence": round(conf, 2),
                         "bounding_box": bbox,
-                        "timestamp": frame_count / fps,
-                        "frame": frame_count
+                        "timestamp": current_time
                     })
 
-            for detection in current_detections:
-                matched_worker = None
-                max_iou = 0
+            frame_count += 1
+
+        video.release()
+        os.remove(video_path)
+        
+        # Process violation history to track workers and persistent violations
+        workers = []
+        for v in violation_history:
+            # Find matching worker
+            matched_worker = None
+            max_iou = 0
+            
+            for worker in workers:
+                iou = calculate_iou(v["bounding_box"], worker["bbox"])
+                if iou > max_iou and iou > CONFIG["IOU_THRESHOLD"]:
+                    max_iou = iou
+                    matched_worker = worker
+
+            if matched_worker:
+                # Update worker's violation history
+                matched_worker["violations"].append(v)
+                matched_worker["bbox"] = v["bounding_box"]
+                matched_worker["last_seen"] = v["timestamp"]
+                v["worker_id"] = matched_worker["id"]
+            else:
+                # New worker
+                worker_id = len(workers) + 1
+                workers.append({
+                    "id": worker_id,
+                    "bbox": v["bounding_box"],
+                    "violations": [v],
+                    "first_seen": v["timestamp"],
+                    "last_seen": v["timestamp"]
+                })
+                v["worker_id"] = worker_id
+
+        # Filter violations to only include those that persist for minimum duration
+        final_violations = []
+        for worker in workers:
+            # Group violations by type
+            violations_by_type = {}
+            for v in worker["violations"]:
+                if v["violation"] not in violations_by_type:
+                    violations_by_type[v["violation"]] = []
+                violations_by_type[v["violation"]].append(v)
+            
+            # Check each violation type for persistence
+            for violation_type, v_list in violations_by_type.items():
+                if len(v_list) < 2:
+                    continue  # Need multiple detections to check duration
                 
-                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:
-                    if detection["violation"] not in matched_worker["violations"]:
-                        logger.info(f"New violation for worker {matched_worker['id']}: {detection['violation']}")
-                        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"]
-                        })
+                duration = max(v["timestamp"] for v in v_list) - min(v["timestamp"] for v in v_list)
+                if duration >= CONFIG["MIN_VIOLATION_DURATION"]:
+                    # Take the highest confidence detection
+                    best_detection = max(v_list, key=lambda x: x["confidence"])
+                    final_violations.append(best_detection)
+                    
+                    # Capture snapshot if not already taken
+                    if not snapshot_taken[violation_type]:
+                        # We need to get the frame for this violation
+                        cap = cv2.VideoCapture(video_path)
+                        cap.set(cv2.CAP_PROP_POS_FRAMES, best_detection["frame"])
+                        ret, snapshot_frame = cap.read()
+                        cap.release()
                         
-                        if not snapshot_taken[detection["violation"]]:
-                            snapshot_filename = f"{detection['violation']}_{frame_count}.jpg"
+                        if ret:
+                            snapshot_filename = f"{violation_type}_{best_detection['frame']}.jpg"
                             snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], snapshot_filename)
-                            cv2.imwrite(snapshot_path, frame)
-                            snapshot_taken[detection["violation"]] = True
+                            cv2.imwrite(snapshot_path, snapshot_frame)
                             snapshots.append({
-                                "violation": detection["violation"],
-                                "frame": frame_count,
+                                "violation": violation_type,
+                                "frame": best_detection["frame"],
                                 "snapshot_path": snapshot_path,
-                                "snapshot_base64": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}",
-                                "timestamp": detection["timestamp"],
-                                "confidence": detection["confidence"]
+                                "snapshot_base64": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}"
                             })
-                    
-                    matched_worker["bbox"] = detection["bounding_box"]
-                    matched_worker["last_frame"] = frame_count
-                else:
-                    worker_id = len(workers) + 1
-                    logger.info(f"New worker {worker_id} with violation: {detection['violation']}")
-                    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
-                    })
-                    
-                    if not snapshot_taken[detection["violation"]]:
-                        snapshot_filename = f"{detection['violation']}_{frame_count}.jpg"
-                        snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], snapshot_filename)
-                        cv2.imwrite(snapshot_path, frame)
-                        snapshot_taken[detection["violation"]] = True
-                        snapshots.append({
-                            "violation": detection["violation"],
-                            "frame": frame_count,
-                            "snapshot_path": snapshot_path,
-                            "snapshot_base64": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}",
-                            "timestamp": detection["timestamp"],
-                            "confidence": detection["confidence"]
-                        })
-
-            active_workers = [w for w in workers if frame_count - w["last_frame"] < CONFIG["FRAME_SKIP"] * 5]
-            if len(active_workers) != len(workers):
-                logger.info(f"Cleaned up {len(workers) - len(active_workers)} inactive workers")
-                workers = active_workers
-                
-            frame_count += 1
+                            snapshot_taken[violation_type] = True
 
-        video.release()
-        os.remove(video_path)
-        
-        violation_types = {}
-        for v in violations:
-            violation_types[v["violation"]] = violation_types.get(v["violation"], 0) + 1
-        
-        logger.info(f"Detection complete. Found violations: {violation_types}")
-
-        if not violations:
-            logger.info("No violations detected")
+        # Final processing
+        if not final_violations:
+            logger.info("No persistent violations detected")
             return {
                 "violations": [],
                 "snapshots": [],
-                "raw_detections": raw_detections,
                 "score": 100,
                 "salesforce_record_id": None,
                 "violation_details_url": "",
                 "message": "No violations detected in the video."
             }
 
-        score = calculate_safety_score(violations)
-        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)
+        score = calculate_safety_score(final_violations)
+        pdf_path, pdf_url, pdf_file = generate_violation_pdf(final_violations, score)
+        report_id, final_pdf_url = push_report_to_salesforce(final_violations, score, pdf_path, pdf_file)
 
         return {
-            "violations": violations,
+            "violations": final_violations,
             "snapshots": snapshots,
-            "raw_detections": raw_detections,
             "score": score,
             "salesforce_record_id": report_id,
             "violation_details_url": final_pdf_url,
@@ -434,18 +446,20 @@ def process_video(video_data):
         return {
             "violations": [],
             "snapshots": [],
-            "raw_detections": [],
             "score": 100,
             "salesforce_record_id": None,
             "violation_details_url": "",
             "message": f"Error processing video: {e}"
         }
 
+# ==========================
+# Gradio Interface (unchanged)
+# ==========================
 def gradio_interface(video_file):
     if not video_file:
-        return "No file uploaded.", "", "No file uploaded.", "", "", [], "No raw detections."
+        return "No file uploaded.", "", "No file uploaded.", "", ""
     try:
-        yield "Processing video... please wait.", "", "", "", "", [], "Processing..."
+        yield "Processing video... please wait.", "", "", "", ""
 
         with open(video_file, "rb") as f:
             video_data = f.read()
@@ -453,71 +467,55 @@ def gradio_interface(video_file):
         result = process_video(video_data)
 
         if result.get("message"):
-            yield result["message"], "", "", "", "", [], "Error in processing."
+            yield result["message"], "", "", "", ""
             return
 
         violation_table = "No violations detected."
         if result["violations"]:
-            header = "| Violation              | Timestamp (s) | Confidence | Worker ID | Frame |\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['worker_id']}        | {v['frame']} |"
+                row = f"| {display_name:<22} | {v['timestamp']:.2f}        | {v['confidence']:.2f}      | {v['worker_id']} |"
                 rows.append(row)
             violation_table = header + separator + "\n".join(rows)
 
         snapshots_text = "No snapshots captured."
-        snapshot_images = []
         if result["snapshots"]:
             violation_name_map = CONFIG["DISPLAY_NAMES"]
             snapshots_text = "\n".join(
-                f"- Snapshot for {violation_name_map.get(s['violation'], s['violation'])} at {s['timestamp']:.2f}s (Frame {s['frame']}, Confidence: {s['confidence']:.2f}): ![]({s['snapshot_base64']})"
+                f"- Snapshot for {violation_name_map.get(s['violation'], s['violation'])} at frame {s['frame']}: ![]({s['snapshot_base64']})"
                 for s in result["snapshots"]
             )
-            snapshot_images = [s["snapshot_base64"] for s in result["snapshots"]]
-
-        raw_detections_text = "No raw detections logged."
-        if result["raw_detections"]:
-            header = "| Frame | Timestamp (s) | Class | Label          | Confidence |\n"
-            separator = "|-------|---------------|-------|----------------|------------|\n"
-            rows = []
-            for d in result["raw_detections"]:
-                row = f"| {d['frame']:<5} | {d['timestamp']:.2f}        | {d['class']:<5} | {d['label']:<14} | {d['confidence']:.2f}      |"
-                rows.append(row)
-            raw_detections_text = header + separator + "\n".join(rows)
 
         yield (
             violation_table,
             f"Safety Score: {result['score']}%",
             snapshots_text,
             f"Salesforce Record ID: {result['salesforce_record_id'] or 'N/A'}",
-            result["violation_details_url"] or "N/A",
-            snapshot_images,
-            raw_detections_text
+            result["violation_details_url"] or "N/A"
         )
     except Exception as e:
         logger.error(f"Error in Gradio interface: {e}", exc_info=True)
-        yield f"Error: {str(e)}", "", "Error in processing.", "", "", [], "Error in processing."
+        yield f"Error: {str(e)}", "", "Error in processing.", "", ""
 
 interface = gr.Interface(
     fn=gradio_interface,
     inputs=gr.Video(label="Upload Site Video"),
-    outputs=[
+    outputs=[ 
         gr.Markdown(label="Detected Safety Violations"),
         gr.Textbox(label="Compliance Score"),
         gr.Markdown(label="Snapshots"),
         gr.Textbox(label="Salesforce Record ID"),
-        gr.Textbox(label="Violation Details URL"),
-        gr.Gallery(label="Violation Snapshots"),
-        gr.Markdown(label="Raw Detections (Debug)")
+        gr.Textbox(label="Violation Details URL")
     ],
     title="Worksite Safety Violation Analyzer",
-    description="Upload site videos to detect safety violations (No Helmet, No Harness, Unsafe Posture). Non-violations are ignored.",
+    description="Upload site videos to detect safety violations (No Helmet, No Harness, Unsafe Posture, Unsafe Zone, Improper Tool Use). Non-violations are ignored.",
     allow_flagging="never"
 )
 
 if __name__ == "__main__":
-    logger.info("Launching Safety Analyzer App...")
+    logger.info("Launching Enhanced Safety Analyzer App...")
     interface.launch()