Update app.py

app.py

@@ -51,22 +51,22 @@ CONFIG = {
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo2/resolve/main/static/output/",
     "FRAME_SKIP": {
-        "no_helmet": 2,  # Reduced to process more frames
+        "no_helmet": 2,
         "no_harness": 1,
         "unsafe_posture": 1,
         "unsafe_zone": 1,
         "improper_tool_use": 1
     },
     "CONFIDENCE_THRESHOLDS": {
-        "no_helmet": 0.5,  # Increased to reduce false positives
-        "no_harness": 0.15,  # Lowered to improve detection
+        "no_helmet": 0.5,
+        "no_harness": 0.15,
         "unsafe_posture": 0.15,
         "unsafe_zone": 0.15,
         "improper_tool_use": 0.15
     },
     "IOU_THRESHOLD": 0.4,
     "MIN_VIOLATION_FRAMES": 3,
-    "HELMET_CONFIDENCE_THRESHOLD": 0.7  # Require high confidence for No Helmet violation
+    "HELMET_CONFIDENCE_THRESHOLD": 0.7
 }
 
 # Setup logging
@@ -493,38 +493,39 @@ def process_video(video_data):
 
         if not violations:
             logger.info("No persistent violations detected")
-            return {
-                "violations": [],
-                "snapshots": [],
-                "score": 100,
-                "salesforce_record_id": None,
-                "violation_details_url": "",
-                "message": "No violations detected in the video."
-            }
+            yield "No violations detected in the video.", "Safety Score: 100%", "No snapshots captured.", "N/A", "N/A"
+            return
 
         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)
 
+        violation_table = "| Violation              | Timestamp (s) | Confidence | Worker ID |\n"
+        violation_table += "|------------------------|---------------|------------|-----------|\n"
+        for v in violations:
+            display_name = CONFIG["DISPLAY_NAMES"].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')} |\n"
+            violation_table += row
+
+        snapshots_text = "No snapshots captured."
+        if snapshots:
+            violation_name_map = CONFIG["DISPLAY_NAMES"]
+            snapshots_text = "\n".join(
+                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 snapshots
+            )
+
         logger.info(f"Processing complete: {len(violations)} violations detected, score: {score}%")
-        return {
-            "violations": violations,
-            "snapshots": snapshots,
-            "score": score,
-            "salesforce_record_id": report_id,
-            "violation_details_url": final_pdf_url,
-            "message": ""
-        }
+        yield (
+            violation_table,
+            f"Safety Score: {score}%",
+            snapshots_text,
+            f"Salesforce Record ID: {report_id or 'N/A'}",
+            final_pdf_url or "N/A"
+        )
     except Exception as e:
         logger.error(f"Error processing video: {e}", exc_info=True)
-        return {
-            "violations": [],
-            "snapshots": [],
-            "score": 100,
-            "salesforce_record_id": None,
-            "violation_details_url": "",
-            "message": f"Error processing video: {e}"
-        }
+        yield f"Error processing video: {e}", "", "", "", ""
 
 # ==========================
 # Gradio Interface
@@ -536,43 +537,9 @@ def gradio_interface(video_file):
         with open(video_file, "rb") as f:
             video_data = f.read()
 
-        # Use generator to yield progress updates
-        for status, violations_table, score, snapshots_text, record_id, details_url in process_video(video_data):
-            if status.startswith("Processing video"):
-                yield status, "", "", "", ""
-                continue
-
-            if status.get("message"):
-                yield status["message"], "", "", "", ""
-                return
-
-            violation_table = "No violations detected."
-            if status["violations"]:
-                header = "| Violation              | Timestamp (s) | Confidence | Worker ID |\n"
-                separator = "|------------------------|---------------|------------|-----------|\n"
-                rows = []
-                violation_name_map = CONFIG["DISPLAY_NAMES"]
-                for v in status["violations"]:
-                    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)
-
-            snapshots_text = "No snapshots captured."
-            if status["snapshots"]:
-                violation_name_map = CONFIG["DISPLAY_NAMES"]
-                snapshots_text = "\n".join(
-                    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 status["snapshots"]
-                )
-
-            yield (
-                violation_table,
-                f"Safety Score: {status['score']}%",
-                snapshots_text,
-                f"Salesforce Record ID: {status['salesforce_record_id'] or 'N/A'}",
-                status["violation_details_url"] or "N/A"
-            )
+        # Use generator to yield updates
+        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.", "", ""
@@ -580,7 +547,7 @@ def gradio_interface(video_file):
 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"),