Update app.py

app.py

@@ -1,12 +1,10 @@
 import os
 import cv2
+import gradio as gr
 import torch
-import time
-from flask import Flask, request, jsonify
+import numpy as np
 from ultralytics import YOLO
-
-# Flask app initialization
-app = Flask(__name__)
+import time
 
 # ==========================
 # Configuration
@@ -29,7 +27,7 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"✅ Using device: {device}")
 
 # ==========================
-# Load Model
+# Load Model (Use YOLOv8n for Faster Inference)
 # ==========================
 selected_model = MODEL_PATH if os.path.isfile(MODEL_PATH) else FALLBACK_MODEL
 model = YOLO(selected_model)
@@ -37,9 +35,9 @@ model = YOLO(selected_model)
 # ==========================
 # Video Processing with Optimizations
 # ==========================
-def process_video(video_file, frame_skip=5, max_frames=100):
+def process_video(video_path, frame_skip=5, max_frames=100):
     try:
-        video = cv2.VideoCapture(video_file)
+        video = cv2.VideoCapture(video_path)
         if not video.isOpened():
             raise ValueError("Could not open video file.")
 
@@ -121,40 +119,23 @@ def generate_pdf_report(violations, score):
     return pdf_url
 
 # ==========================
-# Endpoint for Hugging Face Model Inference
+# Gradio Interface
 # ==========================
-@app.route('/process_video', methods=['POST'])
-def process_video_endpoint():
-    try:
-        # Get the video file from the request
-        if 'video' not in request.files:
-            return jsonify({'error': 'No video file provided'}), 400
-
-        video_file = request.files['video']
-        if not video_file:
-            return jsonify({'error': 'No video file provided'}), 400
+def gradio_interface(video_file):
+    if not video_file:
+        return "Please upload a video file.", ""
 
-        # Save the uploaded video temporarily
-        video_path = os.path.join("temp_video", video_file.filename)
-        video_file.save(video_path)
+    violations, score, pdf_url = process_video(video_file)
+    return violations, f"Safety Score: {score}%", pdf_url
 
-        # Process the video using the model
-        violations, score, pdf_url = process_video(video_path)
-
-        if not violations:
-            return jsonify({'error': 'Error processing video'}), 500
-
-        # Return the violations, safety score, and PDF report URL
-        return jsonify({
-            'violations': violations,
-            'score': score,
-            'pdf_report_url': pdf_url
-        })
-
-    except Exception as e:
-        print(f"❌ Error: {e}")
-        return jsonify({'error': str(e)}), 500
+interface = gr.Interface(
+    fn=gradio_interface,
+    inputs=gr.Video(label="Upload Site Video"),
+    outputs=[gr.JSON(label="Detected Safety Violations"), gr.Textbox(label="Compliance Score"), gr.Textbox(label="PDF Report URL")],
+    title="Worksite Safety Violation Analyzer",
+    description="Upload short site videos to detect safety violations (e.g., no helmet, no harness, unsafe posture)."
+)
 
 if __name__ == "__main__":
-    # Run the Flask app
-    app.run(host="0.0.0.0", port=5000)
+    print("🚀 Launching Safety Analyzer App...")
+    interface.launch()