Update app.py

app.py

@@ -1,10 +1,12 @@
 import os
 import cv2
-import gradio as gr
 import torch
-import numpy as np
-from ultralytics import YOLO
 import time
+from flask import Flask, request, jsonify
+from ultralytics import YOLO
+
+# Flask app initialization
+app = Flask(__name__)
 
 # ==========================
 # Configuration
@@ -27,7 +29,7 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"✅ Using device: {device}")
 
 # ==========================
-# Load Model (Use YOLOv8n for Faster Inference)
+# Load Model
 # ==========================
 selected_model = MODEL_PATH if os.path.isfile(MODEL_PATH) else FALLBACK_MODEL
 model = YOLO(selected_model)
@@ -35,9 +37,9 @@ model = YOLO(selected_model)
 # ==========================
 # Video Processing with Optimizations
 # ==========================
-def process_video(video_path, frame_skip=5, max_frames=100):
+def process_video(video_file, frame_skip=5, max_frames=100):
     try:
-        video = cv2.VideoCapture(video_path)
+        video = cv2.VideoCapture(video_file)
         if not video.isOpened():
             raise ValueError("Could not open video file.")
 
@@ -119,23 +121,40 @@ def generate_pdf_report(violations, score):
     return pdf_url
 
 # ==========================
-# Gradio Interface
+# Endpoint for Hugging Face Model Inference
 # ==========================
-def gradio_interface(video_file):
-    if not video_file:
-        return "Please upload a video file.", ""
+@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
 
-    violations, score, pdf_url = process_video(video_file)
-    return violations, f"Safety Score: {score}%", pdf_url
+        # Save the uploaded video temporarily
+        video_path = os.path.join("temp_video", video_file.filename)
+        video_file.save(video_path)
 
-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)."
-)
+        # 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
 
 if __name__ == "__main__":
-    print("🚀 Launching Safety Analyzer App...")
-    interface.launch()
+    # Run the Flask app
+    app.run(host="0.0.0.0", port=5000)