Update app.py

app.py

@@ -18,7 +18,7 @@ from retrying import retry
 # Configuration
 # ==========================
 CONFIG = {
-    "MODEL_PATH": "yolov8n.pt",  # Force lightweight Nano model
+    "MODEL_PATH": "yolov8n.pt",  # Lightweight model, ensure trained for violations
     "OUTPUT_DIR": "static/output",
     "VIOLATION_LABELS": {
         0: "no_helmet",
@@ -33,8 +33,9 @@ CONFIG = {
         "domain": "login"
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo1/resolve/main/static/output/",
-    "FRAME_SKIP": 15,  # Increased to reduce frames processed
-    "MAX_PROCESSING_TIME": 25  # Cap video processing at 25s to leave time for reporting
+    "FRAME_SKIP": 15,  # Process every 15th frame
+    "MAX_PROCESSING_TIME": 25,  # Cap video processing at 25s
+    "CONFIDENCE_THRESHOLD": 0.5  # Minimum confidence for violation detection
 }
 
 # Setup logging
@@ -57,6 +58,7 @@ def load_model():
     try:
         model = YOLO(CONFIG["MODEL_PATH"]).to(device)
         logger.info(f"Model loaded: {CONFIG['MODEL_PATH']}")
+        logger.warning("Ensure yolov8n.pt is trained to detect only 'no_helmet', 'no_harness', 'unsafe_posture'. Use custom yolov8_safety.pt if needed.")
         return model
     except Exception as e:
         logger.error(f"Failed to load model: {e}")
@@ -250,8 +252,16 @@ def process_video(video_data):
                 for box in result.boxes:
                     cls, conf = int(box.cls), float(box.conf)
                     label = CONFIG["VIOLATION_LABELS"].get(cls, f"class_{cls}")
+                    # Log unexpected classes for debugging
+                    if label not in ["no_helmet", "no_harness", "unsafe_posture", "unsafe_zone"]:
+                        logger.warning(f"Unexpected class detected: {label} (cls: {cls}, conf: {conf})")
+                        continue
+                    # Only process specified violations with sufficient confidence
                     if label not in ["no_helmet", "no_harness", "unsafe_posture"]:
                         continue
+                    if conf < CONFIG["CONFIDENCE_THRESHOLD"]:
+                        logger.info(f"Skipping low-confidence detection: {label} (conf: {conf})")
+                        continue
                     if label in seen_violations:
                         continue
                     seen_violations.add(label)