Update app.py

app.py

@@ -18,8 +18,8 @@ from retrying import retry
 # Configuration
 # ==========================
 CONFIG = {
-    "MODEL_PATH": "yolov8_safety.pt",
-    "FALLBACK_MODEL": "yolov8n.pt",  # updated key to match first code
+    "MODEL_PATH": "yolov8_safety.pt",  # Make sure this file exists in your directory
+    "FALLBACK_MODEL": "yolov8n.pt",  # Fallback model for testing
     "OUTPUT_DIR": "static/output",
     "VIOLATION_LABELS": {
         0: "no_helmet",
@@ -54,11 +54,18 @@ logger.info(f"Using device: {device}")
 
 def load_model():
     try:
-        model_path = CONFIG["MODEL_PATH"] if os.path.isfile(CONFIG["MODEL_PATH"]) else CONFIG["FALLBACK_MODEL"]
-        if model_path == CONFIG["FALLBACK_MODEL"]:
+        # 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)
         model = YOLO(model_path).to(device)
-        logger.info(f"Model loaded: {model_path}")
         return model
     except Exception as e:
         logger.error(f"Failed to load model: {e}")
@@ -309,3 +316,68 @@ def process_video(video_data):
             "violation_details_url": "",
             "message": f"Error processing video: {e}"
         }
+
+def gradio_interface(video_file):
+    if not video_file:
+        return "No file uploaded.", "", "No file uploaded.", "", ""
+    try:
+        yield "Processing video... please wait.", "", "", "", ""
+
+        with open(video_file, "rb") as f:
+            video_data = f.read()
+
+        result = process_video(video_data)
+
+        if result.get("message"):
+            yield result["message"], "", "", "", ""
+            return
+
+        violation_table = "No violations detected."
+        if result["violations"]:
+            header = "| Violation              | Timestamp (s) | Confidence | Bounding Box             |\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']} |"
+                rows.append(row)
+            violation_table = header + separator + "\n".join(rows)
+
+        snapshots_text = "No snapshots captured."
+        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']})"
+                for s in result["snapshots"]
+            )
+
+        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"
+        )
+    except Exception as e:
+        logger.error(f"Error in Gradio interface: {e}", exc_info=True)
+        yield f"Error: {str(e)}", "", "Error in processing.", "", ""
+
+interface = gr.Interface(
+    fn=gradio_interface,
+    inputs=gr.Video(label="Upload Site Video"),
+    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")
+    ],
+    title="Worksite Safety Violation Analyzer",
+    description="Upload site videos to detect safety violations (No Helmet Violation, No Harness Violation, Unsafe Posture Violation). Non-violations are ignored.",
+    allow_flagging="never"
+)
+
+if __name__ == "__main__":
+    logger.info("Launching Safety Analyzer App...")
+    interface.launch()