Fixed mismatch of the helmet

app.py

@@ -185,7 +185,7 @@ with gr.Blocks(title="Small Object Detection") as app:
                 ppe_threshold_slider = gr.Slider(
                     minimum=0.05,
                     maximum=0.95,
-                    value=0.3,
+                    value=0.25,
                     step=0.05,
                     label="PPE detection threshold",
                 )

ppe_compliance.py

@@ -34,6 +34,18 @@ DEVICE = "cpu"
 
 ALL_PPE = ["goggles", "helmet", "mask", "shoes", "vest", "glove"]
 
+# Expected vertical position of each item within its wearer's box (0 = top of the
+# person box, 1 = bottom). Used ONLY to disambiguate which of several overlapping
+# persons an item belongs to — never to reject an item outright.
+PPE_VFRAC = {
+    "helmet":  0.10,
+    "goggles": 0.13,
+    "mask":    0.20,
+    "vest":    0.45,
+    "glove":   0.62,
+    "shoes":   0.92,
+}
+
 # Person detector defaults to the same model the "Detect & Classify" tab uses.
 DEFAULT_PERSON_MODEL = "medium-obj2coco"
 
@@ -119,7 +131,7 @@ def _color_of(name):
 def run_ppe_compliance(
     image,
     person_threshold=0.75,
-    ppe_threshold=0.3,
+    ppe_threshold=0.25,
     assoc=0.5,
     person_model=DEFAULT_PERSON_MODEL,
     min_side=960,
@@ -151,17 +163,12 @@ def run_ppe_compliance(
     persons.sort(key=lambda ps: -ps[0])
     persons = persons[:MAX_PERSON_CROPS]
 
-    # 2+3) PPE detector (fine-tuned D-FINE-M), run on EACH person crop so small PPE
-    # (vests, distant helmets) isn't lost to the 640x640 full-frame resize, then
-    # attribute the detections to THAT person. We credit a detection to the person
-    # whose crop produced it (gated by containment in their own box),
-    # rather than re-assigning globally to whichever overlapping neighbour contains
-    # it most — that global step mislabelled a worker's own vest to a denser
-    # neighbour in crowded frames, so the vest was drawn but the wearer showed
-    # "no vest". Boxes are mapped back to full coords; the draw list is de-duped.
+    # 2) PPE detector (fine-tuned D-FINE-M), run on EACH person crop so small PPE
+    # (vests, distant helmets) isn't lost to the 640x640 full-frame resize. Boxes
+    # are mapped back to full-image coords; the union is de-duplicated across crops.
     people = [{"score": conf, "box": pb, "present": {}} for conf, pb in persons]
-    ppe = []  # (name, score, full_box) — for drawing only
-    for i, (conf, pb) in enumerate(persons):
+    ppe = []  # (name, score, full_box)
+    for conf, pb in persons:
         x1, y1, x2, y2 = pb
         pw, ph = x2 - x1, y2 - y1
         cx1 = max(0, int(x1 - PERSON_CROP_PAD * pw))
@@ -172,15 +179,31 @@ def run_ppe_compliance(
             continue
         crop = im.crop((cx1, cy1, cx2, cy2))
         for name, s, b in detect_ppe_boxes(crop, threshold=ppe_threshold):
-            box = [b[0] + cx1, b[1] + cy1, b[2] + cx1, b[3] + cy1]
-            ppe.append((name, s, box))
-            # Attribute to this crop's person: must lie mostly inside their box.
+            ppe.append((name, s, [b[0] + cx1, b[1] + cy1, b[2] + cx1, b[3] + cy1]))
+    ppe = _dedup_ppe(ppe)
+
+    # 3) Attribute each PPE item to EXACTLY ONE person. Among the persons whose box
+    # contains the item (>= assoc), pick the one where it sits in the anatomically
+    # expected place (helmet near the top, shoes near the bottom, ...). This stops a
+    # helmet that lies in the MIDDLE of a tall overlapping neighbour from being
+    # credited to them when it's really at the TOP (head) of the person beside them.
+    # It's a tie-break among containers, not a hard reject, so a uniquely-contained
+    # item is still credited regardless of pose (no false "missing" on bent workers).
+    for name, s, box in ppe:
+        tgt = PPE_VFRAC.get(name, 0.5)
+        cy = 0.5 * (box[1] + box[3])
+        best_i, best_d = -1, None
+        for i, p in enumerate(people):
+            pb = p["box"]
             if _contain_frac(box, pb) < assoc:
                 continue
-            cur = people[i]["present"].get(name)
+            d = abs((cy - pb[1]) / max(1e-6, pb[3] - pb[1]) - tgt)
+            if best_d is None or d < best_d:
+                best_i, best_d = i, d
+        if best_i >= 0:
+            cur = people[best_i]["present"].get(name)
             if cur is None or s > cur[0]:
-                people[i]["present"][name] = (s, box)
-    ppe = _dedup_ppe(ppe)
+                people[best_i]["present"][name] = (s, box)
 
     # 4) Verdicts + report.
     verdicts = []