Update app.py

app.py

@@ -7,13 +7,14 @@ import numpy as np
 try:
     from ultralytics import YOLO
 except ImportError as e:
-    print("❌ Ultralytics not installed.")
+    print("❌ Ultralytics not installed. Run: pip install ultralytics")
     raise
 
 # ==========================
 # Configuration
 # ==========================
 DEFAULT_MODEL_PATH = "models/yolov8_safety.pt"
+FALLBACK_MODEL = "yolov8n.pt"  # Use nano model if custom one is missing
 MODEL_PATH = os.getenv("SAFETY_MODEL_PATH", DEFAULT_MODEL_PATH)
 
 VIOLATION_LABELS = {
@@ -36,15 +37,16 @@ except Exception as e:
 # ==========================
 # Load Model
 # ==========================
-if not os.path.isfile(MODEL_PATH):
-    raise FileNotFoundError(
-        f"❌ ERROR: Model file '{MODEL_PATH}' not found.\n"
-        f"👉 Please upload it to the 'models/' folder or set SAFETY_MODEL_PATH env variable."
-    )
+if os.path.isfile(MODEL_PATH):
+    selected_model = MODEL_PATH
+    print(f"✅ Found model at: {selected_model}")
+else:
+    print(f"⚠️ Model file '{MODEL_PATH}' not found. Falling back to: {FALLBACK_MODEL}")
+    selected_model = FALLBACK_MODEL
 
 try:
-    model = YOLO(MODEL_PATH)
-    print(f"✅ Model loaded from {MODEL_PATH}")
+    model = YOLO(selected_model)
+    print(f"✅ Model loaded: {selected_model}")
 except Exception as e:
     print(f"❌ Failed to load model: {e}")
     raise
@@ -74,13 +76,13 @@ def process_video(video_path):
                     conf = float(box.conf)
                     xywh = box.xywh.cpu().numpy()[0]
 
-                    if cls in VIOLATION_LABELS:
-                        violations.append({
-                            "frame": frame_count,
-                            "violation": VIOLATION_LABELS[cls],
-                            "confidence": round(conf, 2),
-                            "bounding_box": [round(x, 2) for x in xywh]
-                        })
+                    label = VIOLATION_LABELS.get(cls, f"class_{cls}")
+                    violations.append({
+                        "frame": frame_count,
+                        "violation": label,
+                        "confidence": round(conf, 2),
+                        "bounding_box": [round(x, 2) for x in xywh]
+                    })
 
             frame_count += 1