Update app.py

app.py

@@ -7,11 +7,15 @@ import numpy as np
 try:
     from ultralytics import YOLO
 except ImportError as e:
-    print(f"Error importing ultralytics: {e}")
+    print("❌ Ultralytics not installed.")
     raise
 
-# ========== Configuration ==========
-MODEL_PATH = "models/yolov8_safety.pt"  # Your custom safety model
+# ==========================
+# Configuration
+# ==========================
+DEFAULT_MODEL_PATH = "models/yolov8_safety.pt"
+MODEL_PATH = os.getenv("SAFETY_MODEL_PATH", DEFAULT_MODEL_PATH)
+
 VIOLATION_LABELS = {
     0: "no_helmet",
     1: "no_harness",
@@ -19,41 +23,49 @@ VIOLATION_LABELS = {
     3: "unsafe_zone"
 }
 
-# ========== Device Setup ==========
+# ==========================
+# Device Setup
+# ==========================
 try:
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    print(f"Using device: {device}")
+    print(f"✅ Using device: {device}")
 except Exception as e:
-    print(f"Device error: {e}")
+    print(f"⚠️ Error setting device: {e}")
     device = torch.device("cpu")
 
-# ========== Load Model ==========
+# ==========================
+# 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.")
+        f"❌ ERROR: Model file '{MODEL_PATH}' not found.\n"
+        f"👉 Please upload it to the 'models/' folder or set SAFETY_MODEL_PATH env variable."
+    )
 
 try:
     model = YOLO(MODEL_PATH)
+    print(f"✅ Model loaded from {MODEL_PATH}")
 except Exception as e:
-    print(f"Error loading YOLO model: {e}")
+    print(f"❌ Failed to load model: {e}")
     raise
 
-# ========== Core Logic ==========
+# ==========================
+# Video Processing
+# ==========================
 def process_video(video_path):
     try:
         video = cv2.VideoCapture(video_path)
         if not video.isOpened():
             raise ValueError("Could not open video file.")
 
-        violations = []
         frame_count = 0
+        violations = []
 
         while True:
             ret, frame = video.read()
             if not ret:
                 break
 
-            # YOLOv8 inference
             results = model(frame, device=device)
 
             for result in results:
@@ -73,14 +85,16 @@ def process_video(video_path):
             frame_count += 1
 
         video.release()
-        safety_score = calculate_safety_score(violations)
-        return violations, safety_score
+        score = calculate_safety_score(violations)
+        return violations, score
 
     except Exception as e:
-        print(f"Error processing video: {e}")
+        print(f"❌ Error processing video: {e}")
         return [], f"Error: {e}"
 
-# ========== Score Calculation ==========
+# ==========================
+# Safety Score
+# ==========================
 def calculate_safety_score(violations):
     base_score = 100
     penalties = {
@@ -93,7 +107,9 @@ def calculate_safety_score(violations):
         base_score -= penalties.get(v["violation"], 0)
     return max(base_score, 0)
 
-# ========== Gradio Interface ==========
+# ==========================
+# Gradio Interface
+# ==========================
 def gradio_interface(video_file):
     if not video_file:
         return "Please upload a video file.", ""
@@ -109,9 +125,9 @@ interface = gr.Interface(
         gr.Textbox(label="Compliance 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."
+    description="Upload short site videos to detect safety violations (e.g., no helmet, no harness, unsafe posture)."
 )
 
 if __name__ == "__main__":
-    print("Launching Safety Analyzer App...")
+    print("🚀 Launching Safety Analyzer App...")
     interface.launch()