Update app.py

app.py

@@ -18,7 +18,7 @@ from retrying import retry
 # Configuration
 # ==========================
 CONFIG = {
-    "MODEL_PATH": "yolov8n.pt",  # Lightweight model, ensure trained for violations
+    "MODEL_PATH": "yolov8n.pt",  # Lightweight model, must be trained for violations only
     "OUTPUT_DIR": "static/output",
     "VIOLATION_LABELS": {
         0: "no_helmet",
@@ -58,7 +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.")
+        logger.warning("Ensure yolov8n.pt is trained to detect ONLY 'no_helmet', 'no_harness', 'unsafe_posture'. Replace with custom-trained yolov8_safety.pt if positive cases are detected.")
         return model
     except Exception as e:
         logger.error(f"Failed to load model: {e}")
@@ -252,13 +252,14 @@ 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
+                    # Log and skip any unexpected classes
                     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
+                    # Only process specified violations
                     if label not in ["no_helmet", "no_harness", "unsafe_posture"]:
                         continue
+                    # Apply confidence threshold
                     if conf < CONFIG["CONFIDENCE_THRESHOLD"]:
                         logger.info(f"Skipping low-confidence detection: {label} (conf: {conf})")
                         continue
@@ -376,7 +377,7 @@ interface = gr.Interface(
         gr.Textbox(label="Violation Details URL")
     ],
     title="Worksite Safety Violation Analyzer",
-    description="Upload site videos to detect safety violations (e.g., no helmet, no harness, unsafe posture)."
+    description="Upload site videos to detect safety violations (no helmet, no harness, unsafe posture). Positive cases are ignored."
 )
 
 if __name__ == "__main__":