Update app.py

app.py

@@ -1,65 +1,73 @@
-#import cv2
+import os
+import cv2
 import gradio as gr
-import torch  # Moved torch import to the top
+import torch
+import numpy as np
+
 try:
     from ultralytics import YOLO
 except ImportError as e:
     print(f"Error importing ultralytics: {e}")
-    print("Ensure 'ultralytics' is listed in requirements.txt and installed.")
     raise
-import numpy as np
 
-# Set device for model inference
+# ========== Configuration ==========
+MODEL_PATH = "models/yolov8_safety.pt"  # Your custom safety model
+VIOLATION_LABELS = {
+    0: "no_helmet",
+    1: "no_harness",
+    2: "unsafe_posture",
+    3: "unsafe_zone"
+}
+
+# ========== Device Setup ==========
 try:
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     print(f"Using device: {device}")
 except Exception as e:
-    print(f"Error setting device: {e}")
-    device = torch.device("cpu")  # Fallback to CPU
-    print("Falling back to CPU")
+    print(f"Device error: {e}")
+    device = torch.device("cpu")
+
+# ========== Load Model ==========
+if not os.path.isfile(MODEL_PATH):
+    raise FileNotFoundError(
+        f"🚨 ERROR: Model file '{MODEL_PATH}' not found. Please upload it to the 'models/' folder.")
 
-# Load the YOLOv8 model
 try:
-    model = YOLO("yolov8n.pt")  # Use YOLOv8 nano model
+    model = YOLO(MODEL_PATH)
 except Exception as e:
     print(f"Error loading YOLO model: {e}")
     raise
 
-# Function to process the video file
+# ========== Core Logic ==========
 def process_video(video_path):
     try:
-        # Load the video
         video = cv2.VideoCapture(video_path)
         if not video.isOpened():
             raise ValueError("Could not open video file.")
-       
-        frame_count = 0
+
         violations = []
+        frame_count = 0
 
         while True:
             ret, frame = video.read()
             if not ret:
-                break  # End of video
+                break
 
-            # Run YOLOv8 inference on the frame
+            # YOLOv8 inference
             results = model(frame, device=device)
 
-            # Process detected objects
             for result in results:
-                boxes = result.boxes
-                for box in boxes:
+                for box in result.boxes:
                     cls = int(box.cls)
                     conf = float(box.conf)
                     xywh = box.xywh.cpu().numpy()[0]
 
-                    # Map class IDs to violation types (adjust as needed)
-                    violation_labels = {0: "person", 1: "bicycle", 2: "car"}
-                    if cls in violation_labels:
+                    if cls in VIOLATION_LABELS:
                         violations.append({
                             "frame": frame_count,
-                            "violation": violation_labels.get(cls, "unknown"),
-                            "confidence": conf,
-                            "bounding_box": xywh.tolist()
+                            "violation": VIOLATION_LABELS[cls],
+                            "confidence": round(conf, 2),
+                            "bounding_box": [round(x, 2) for x in xywh]
                         })
 
             frame_count += 1
@@ -67,47 +75,43 @@ def process_video(video_path):
         video.release()
         safety_score = calculate_safety_score(violations)
         return violations, safety_score
+
     except Exception as e:
         print(f"Error processing video: {e}")
         return [], f"Error: {e}"
 
-# Function to calculate safety score
+# ========== Score Calculation ==========
 def calculate_safety_score(violations):
-    total_score = 100
-    violation_penalties = {
-        "person": 20,
-        "bicycle": 15,
-        "car": 30,
-        "unknown": 10
+    base_score = 100
+    penalties = {
+        "no_helmet": 25,
+        "no_harness": 30,
+        "unsafe_posture": 20,
+        "unsafe_zone": 25
     }
-    for violation in violations:
-        total_score -= violation_penalties.get(violation["violation"], 0)
-    return max(total_score, 0)
+    for v in violations:
+        base_score -= penalties.get(v["violation"], 0)
+    return max(base_score, 0)
 
-# Gradio Interface
+# ========== Gradio Interface ==========
 def gradio_interface(video_file):
-    if video_file is None:
+    if not video_file:
         return "Please upload a video file.", ""
-   
-    try:
-        violations, safety_score = process_video(video_file)
-        return violations, f"Safety Score: {safety_score}%"
-    except Exception as e:
-        print(f"Gradio interface error: {e}")
-        return [], f"Error: {e}"
 
-# Define Gradio interface
+    violations, score = process_video(video_file)
+    return violations, f"Safety Score: {score}%"
+
 interface = gr.Interface(
     fn=gradio_interface,
-    inputs=gr.Video(label="Upload Video"),
+    inputs=gr.Video(label="Upload Site Video"),
     outputs=[
-        gr.JSON(label="Detected Violations"),
-        gr.Textbox(label="Safety Score")
+        gr.JSON(label="Detected Safety Violations"),
+        gr.Textbox(label="Compliance Score")
     ],
-    title="Safety Violation Detection",
-    description="Upload a video to detect safety violations and calculate a safety score."
+    title="Worksite Safety Violation Analyzer",
+    description="Upload a short site video to detect safety compliance violations like missing helmets, no harness, and unsafe behavior."
 )
 
 if __name__ == "__main__":
-    print("Launching Gradio interface...")
-    interface.launch()
+    print("Launching Safety Analyzer App...")
+    interface.launch()