Update app.py

app.py

@@ -51,7 +51,7 @@ CONFIG = {
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo2/resolve/main/static/output/",
     "FRAME_SKIP": {
-        "no_helmet": 3,  # Lower skip for frequent violations
+        "no_helmet": 3,
         "no_harness": 2,
         "unsafe_posture": 2,
         "unsafe_zone": 2,
@@ -59,14 +59,14 @@ CONFIG = {
     },
     "MAX_PROCESSING_TIME": 60,
     "CONFIDENCE_THRESHOLDS": {
-        "no_helmet": 0.3,  # Slightly higher to reduce false positives
+        "no_helmet": 0.3,
         "no_harness": 0.2,
         "unsafe_posture": 0.2,
         "unsafe_zone": 0.2,
         "improper_tool_use": 0.2
     },
     "IOU_THRESHOLD": 0.4,
-    "MIN_VIOLATION_FRAMES": 3  # Minimum consecutive frames to confirm violation
+    "MIN_VIOLATION_FRAMES": 3
 }
 
 # Setup logging
@@ -103,9 +103,9 @@ model = load_model()
 def draw_detections(frame, detections):
     """Draw bounding boxes and labels on frame"""
     for det in detections:
-        label = det["violation"]
-        confidence = det["confidence"]
-        x, y, w, h = det["bounding_box"]
+        label = det.get("violation", "Unknown")
+        confidence = det.get("confidence", 0.0)
+        x, y, w, h = det.get("bounding_box", [0, 0, 0, 0])
         
         x1 = int(x - w/2)
         y1 = int(y - h/2)
@@ -178,8 +178,8 @@ def generate_violation_pdf(violations, score):
             c.drawString(1 * inch, y_position, "No violations detected.")
         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']})"
+                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})"
                 c.drawString(1 * inch, y_position, text)
                 y_position -= 0.3 * inch
                 if y_position < 1 * inch:
@@ -228,7 +228,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']})"
+            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})"
             for v in violations
         ) or "No violations detected."
         pdf_url = f"{CONFIG['PUBLIC_URL_BASE']}{os.path.basename(pdf_path)}" if pdf_path else ""
@@ -275,7 +275,7 @@ def calculate_safety_score(violations):
         "unsafe_zone": 35,
         "improper_tool_use": 25
     }
-    total_penalty = sum(penalties.get(v["violation"], 0) for v in violations)
+    total_penalty = sum(penalties.get(v.get("violation", "Unknown"), 0) for v in violations)
     score = 100 - total_penalty
     return max(score, 0)
 
@@ -337,24 +337,32 @@ def process_video(video_data):
                     label = CONFIG["VIOLATION_LABELS"].get(cls, None)
                     
                     if label is None:
+                        logger.warning(f"Unknown class ID {cls} detected, skipping")
                         continue
                         
                     if conf < CONFIG["CONFIDENCE_THRESHOLDS"].get(label, 0.25):
+                        logger.debug(f"Detection {label} with confidence {conf:.2f} below threshold, skipping")
                         continue
 
                     bbox = [round(x, 2) for x in box.xywh.cpu().numpy()[0]]
                     
                     current_detections.append({
                         "frame": frame_count,
-                        "vi臉olation": label,
+                        "violation": label,  # Corrected key
                         "confidence": round(conf, 2),
                         "bounding_box": bbox,
                         "timestamp": current_time
                     })
 
+            logger.debug(f"Frame {frame_count}: Detected {len(current_detections)} violations: {[d['violation'] for d in current_detections]}")
+
             # Process detections and associate with workers
             for detection in current_detections:
-                violation_type = detection["violation"]
+                violation_type = detection.get("violation", None)
+                if violation_type is None:
+                    logger.error(f"Invalid detection, missing 'violation' key: {detection}")
+                    continue
+
                 # Skip No Helmet detection if worker is compliant
                 if violation_type == "no_helmet":
                     matched_worker = None
@@ -366,7 +374,8 @@ def process_video(video_data):
                             matched_worker = worker
                     
                     if matched_worker and matched_worker["id"] in helmet_compliance:
-                        continue  # Skip if worker is known to wear a helmet
+                        logger.debug(f"Worker {matched_worker['id']} has helmet, skipping no_helmet violation")
+                        continue
                 
                 # Find or create worker
                 matched_worker = None
@@ -393,14 +402,16 @@ def process_video(video_data):
                 
                 # Skip if this violation type is already confirmed for this worker
                 if worker_id in confirmed_violations and violation_type in confirmed_violations[worker_id]:
+                    logger.debug(f"Violation {violation_type} already confirmed for worker {worker_id}, skipping")
                     continue
                 
                 detection["worker_id"] = worker_id
                 violation_history[violation_type].append(detection)
                 
                 # Update helmet compliance (simulate by checking if No Helmet confidence is low)
-                if violation_type == "no_helmet" and conf < 0.5:
+                if violation_type == "no_helmet" and detection["confidence"] < 0.5:
                     helmet_compliance[worker_id] = True
+                    logger.debug(f"Worker {worker_id} marked as helmet compliant")
 
             # Clean up old workers
             workers = [w for w in workers if current_time - w["last_seen"] < 5.0]
@@ -413,6 +424,7 @@ def process_video(video_data):
         # Process violation history to confirm persistent violations
         for violation_type, detections in violation_history.items():
             if not detections:
+                logger.info(f"No detections for {violation_type}")
                 continue
                 
             worker_violations = {}
@@ -429,6 +441,7 @@ def process_video(video_data):
                         
                     # Skip No Helmet if worker is compliant
                     if violation_type == "no_helmet" and worker_id in helmet_compliance:
+                        logger.debug(f"Skipping no_helmet for worker {worker_id} due to helmet compliance")
                         continue
                     
                     best_detection = max(worker_dets, key=lambda x: x["confidence"])
@@ -457,6 +470,7 @@ def process_video(video_data):
                                 "snapshot_base64": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}"
                             })
                             snapshot_taken[violation_type] = True
+                            logger.info(f"Snapshot taken for {violation_type} at frame {best_detection['frame']}")
 
         if not violations:
             logger.info("No persistent violations detected")
@@ -473,6 +487,7 @@ 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)
 
+        logger.info(f"Processing complete: {len(violations)} violations detected, score: {score}%")
         return {
             "violations": violations,
             "snapshots": snapshots,
@@ -517,8 +532,8 @@ def gradio_interface(video_file):
             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']} |"
+                display_name = violation_name_map.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')} |"
                 rows.append(row)
             violation_table = header + separator + "\n".join(rows)
 
@@ -526,7 +541,7 @@ def gradio_interface(video_file):
         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 frame {s['frame']}: ![]({s['snapshot_base64']})"
+                f"- Snapshot for {violation_name_map.get(s.get('violation', 'Unknown'), 'Unknown')} at frame {s.get('frame', 0)}: ![]({s.get('snapshot_base64', '')})"
                 for s in result["snapshots"]
             )