Update app.py

app.py

@@ -24,7 +24,12 @@ CONFIG = {
         0: "no_helmet",
         1: "no_harness",
         2: "unsafe_posture",
-        3: "unsafe_zone"  # Ignored in scoring
+        3: "unsafe_zone"  # Ignored in scoring and table
+    },
+    "DISPLAY_NAMES": {  # Mapping for user-friendly violation names
+        "no_helmet": "Missing Helmet",
+        "no_harness": "Missing Harness",
+        "unsafe_posture": "Unsafe Posture"
     },
     "SF_CREDENTIALS": {
         "username": "prashanth1ai@safety.com",
@@ -58,7 +63,7 @@ def load_model():
     try:
         model = YOLO(CONFIG["MODEL_PATH"]).to(device)
         logger.info(f"Model loaded: {CONFIG['MODEL_PATH']}")
-        logger.warning("Ensure yolov8n.pt is trained to detect ONLY 'no_helmet', 'no_harness', 'unsafe_posture'. Replace with custom-trained yolov8_safety.pt if positive cases are detected.")
+        logger.warning("Ensure yolov8n.pt is trained to detect ONLY 'no_helmet', 'no_harness', 'unsafe_posture'. Replace with custom-trained yolov8_safety.pt if unexpected classes are detected.")
         return model
     except Exception as e:
         logger.error(f"Failed to load model: {e}")
@@ -104,7 +109,8 @@ def generate_violation_pdf(violations, score):
         c.drawString(1 * inch, y_position, "Violation Details:")
         y_position -= 0.3 * inch
         for v in violations:
-            text = f"{v['violation']} at {v['timestamp']:.2f}s (Confidence: {v['confidence']})"
+            display_name = CONFIG["DISPLAY_NAMES"].get(v["violation"], v["violation"])
+            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:
@@ -155,7 +161,7 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
     try:
         sf = connect_to_salesforce()
         violations_text = "\n".join(
-            f"{v['violation']} at {v['timestamp']:.2f}s (Confidence: {v['confidence']})"
+            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 ""
@@ -251,13 +257,10 @@ def process_video(video_data):
             for result in results:
                 for box in result.boxes:
                     cls, conf = int(box.cls), float(box.conf)
-                    label = CONFIG["VIOLATION_LABELS"].get(cls, f"class_{cls}")
-                    # Log and skip any unexpected classes
-                    if label not in ["no_helmet", "no_harness", "unsafe_posture", "unsafe_zone"]:
-                        logger.warning(f"Unexpected class detected: {label} (cls: {cls}, conf: {conf})")
-                        continue
+                    label = CONFIG["VIOLATION_LABELS"].get(cls, f"unknown_class_{cls}")
                     # Only process specified violations
                     if label not in ["no_helmet", "no_harness", "unsafe_posture"]:
+                        logger.warning(f"Ignoring detection: {label} (cls: {cls}, conf: {conf}) - not a target violation")
                         continue
                     # Apply confidence threshold
                     if conf < CONFIG["CONFIDENCE_THRESHOLD"]:
@@ -339,19 +342,19 @@ def gradio_interface(video_file):
 
         violation_table = "No violations detected."
         if result["violations"]:
-            header = "| Violation     | Timestamp | Confidence | Bounding Box             | Violation Details       |\n"
-            separator = "|---------------|-----------|------------|--------------------------|-------------------------|\n"
+            header = "| Violation        | Timestamp | Confidence | Bounding Box             | Violation Details       |\n"
+            separator = "|------------------|-----------|------------|--------------------------|-------------------------|\n"
             rows = []
             for v in result["violations"]:
-                violation_name = v["violation"].replace("no_", "").replace("unsafe_", "")
-                row = f"| {violation_name:<13} | {v['timestamp']:.2f}s  | {v['confidence']:.2f}      | {v['bounding_box']} | {result['violation_details_url']} |"
+                display_name = CONFIG["DISPLAY_NAMES"].get(v["violation"], v["violation"])
+                row = f"| {display_name:<16} | {v['timestamp']:.2f}s  | {v['confidence']:.2f}      | {v['bounding_box']} | {result['violation_details_url']} |"
                 rows.append(row)
             violation_table = header + separator + "\n".join(rows)
 
         snapshots_text = "No snapshots captured."
         if result["snapshots"]:
             snapshots_text = "\n".join(
-                f"- Snapshot for {s['violation'].replace('no_', '').replace('unsafe_', '')} at frame {s['frame']}: ![]({s['snapshot_base64']})"
+                f"- Snapshot for {CONFIG['DISPLAY_NAMES'].get(s['violation'], s['violation'])} at frame {s['frame']}: ![]({s['snapshot_base64']})"
                 for s in result["snapshots"]
             )
 
@@ -366,7 +369,7 @@ def gradio_interface(video_file):
         logger.error(f"Error in Gradio interface: {e}")
         return f"Error: {str(e)}", "", "Error in processing.", "", ""
 
-interface = gr.Interface(
+interface = gr.example(
     fn=gradio_interface,
     inputs=gr.Video(label="Upload Site Video"),
     outputs=[
@@ -377,7 +380,7 @@ interface = gr.Interface(
         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). Positive cases are ignored."
+    description="Upload site videos to detect safety violations (Missing Helmet, Missing Harness, Unsafe Posture). Non-violations are ignored."
 )
 
 if __name__ == "__main__":