Update app.py

app.py

@@ -53,7 +53,7 @@ CONFIG = {
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo2/resolve/main/static/output/",
     "CONFIDENCE_THRESHOLDS": {
-        "no_helmet": 0.6,
+        "no_helmet": 0.75,  # Increased for stricter helmet detection
         "no_harness": 0.4,
         "unsafe_posture": 0.4,
         "unsafe_zone": 0.4,
@@ -422,11 +422,13 @@ def process_video(video_data):
 
                     detection["worker_id"] = worker_id
 
-                    # Track helmet violations separately
+                    # Track helmet violations with stricter criteria
                     if detection["violation"] == "no_helmet":
-                        if worker_id not in helmet_violations:
-                            helmet_violations[worker_id] = []
-                        helmet_violations[worker_id].append(detection)
+                        # Only include high-confidence no_helmet detections
+                        if conf >= CONFIG["CONFIDENCE_THRESHOLDS"]["no_helmet"]:
+                            if worker_id not in helmet_violations:
+                                helmet_violations[worker_id] = []
+                            helmet_violations[worker_id].append(detection)
                     else:
                         violations.append(detection)
 
@@ -441,10 +443,11 @@ def process_video(video_data):
         # Confirm helmet violations (require multiple detections)
         for worker_id, detections in helmet_violations.items():
             if len(detections) >= CONFIG["MIN_VIOLATION_FRAMES"]:
+                # Select the detection with the highest confidence
                 best_detection = max(detections, key=lambda x: x["confidence"])
                 violations.append(best_detection)
 
-                # Capture snapshot
+                # Capture snapshot for confirmed no_helmet violation
                 cap = cv2.VideoCapture(video_path)
                 cap.set(cv2.CAP_PROP_POS_FRAMES, best_detection["frame"])
                 ret, snapshot_frame = cap.read()