Spaces:
Sleeping
Sleeping
Download detector.py from ZahraFatima/aigis-captioning: direct link, hf CLI and curl.
- Browser
- Download file 17 kB
-
https://huggingface.co/spaces/ZahraFatima/aigis-captioning/resolve/main/detector.py
- Command line
-
hf download hf://spaces/ZahraFatima/aigis-captioning/detector.py
-
curl -L -o detector.py https://huggingface.co/spaces/ZahraFatima/aigis-captioning/resolve/main/detector.py
17 kB
| """ | |
| Wrap a trained YOLO model into structured, VLM-ready detections. | |
| This is the DETERMINISTIC half of the pipeline. It: | |
| 1. runs YOLO on an image | |
| 2. returns clean structured detections (class, confidence, box) | |
| 3. does simple, auditable PPE-compliance reasoning, PER PERSON: | |
| "this person's head is detected and no helmet is on it" -> helmet missing | |
| The VLM never decides compliance on its own. It is handed these facts, and the app | |
| renders them as a table straight from this module. This keeps the safety-critical | |
| judgement in code that can be logged and reviewed, in line with the AiGiS design | |
| principle. | |
| What is checked comes from the data YAML (sh17.yaml), not from this file: | |
| ppe_rules: | |
| - item: eye_protection # name shown in the PPE table | |
| body_part: face # detector class the item protects | |
| any_of: [glasses, face_guard] # ANY one of these satisfies the rule | |
| level: advisory # "required" or "advisory" | |
| compliance_conf: 0.40 # confidence floor for compliance evidence | |
| The older format is still accepted (every listed item is treated as required): | |
| ppe_pairs: | |
| face: [glasses, face_mask_medical] | |
| How one person's item gets its status: | |
| present a confident PPE box sits on the body part (or is on the person, for | |
| items such as gloves where the gloved hand is not detected as "hands") | |
| missing the body part is confidently detected and no PPE box is on it. | |
| PPE worn by OTHER people never changes this. | |
| uncertain the only evidence is a low-confidence box, or the PPE is on the | |
| person but not clearly over the body part | |
| A body part that is not detected at all produces no finding (not assessed). | |
| """ | |
| from __future__ import annotations | |
| from dataclasses import dataclass, asdict, field | |
| from pathlib import Path | |
| import yaml | |
| from ultralytics import YOLO | |
| class Detection: | |
| cls: str | |
| confidence: float | |
| box_xywhn: list[float] # normalized [cx, cy, w, h] in [0,1] | |
| def to_dict(self) -> dict: | |
| return asdict(self) | |
| class ComplianceFinding: | |
| body_part: str | |
| expected_ppe: str # class name(s) that satisfy the rule, e.g. "glasses or face_guard" | |
| status: str # "present" | "missing" | "uncertain" | |
| detail: str | |
| # --- added fields (older consumers can ignore them) --- | |
| item: str = "" # rule name, e.g. "eye_protection" | |
| level: str = "required" # "required" | "advisory" | |
| required: bool = True | |
| person_id: int | None = None # 1..N, people numbered left to right; None = not tied to a person | |
| confidence: float | None = None # confidence of the box that decided the status | |
| def to_dict(self) -> dict: | |
| return asdict(self) | |
| class DetectionResult: | |
| image: str | |
| detections: list[Detection] = field(default_factory=list) | |
| compliance: list[ComplianceFinding] = field(default_factory=list) | |
| people: int = 0 # number of people the compliance check assessed | |
| def to_dict(self) -> dict: | |
| return { | |
| "image": self.image, | |
| "people": self.people, | |
| "detections": [d.to_dict() for d in self.detections], | |
| "compliance": [c.to_dict() for c in self.compliance], | |
| } | |
| class PpeRule: | |
| item: str | |
| body_part: str | |
| any_of: list[str] | |
| level: str = "required" | |
| # --------------------------------------------------------------------------- | |
| # Geometry | |
| # --------------------------------------------------------------------------- | |
| def _xywhn_to_corners(b: list[float]) -> tuple[float, float, float, float]: | |
| cx, cy, w, h = b | |
| return cx - w / 2, cy - h / 2, cx + w / 2, cy + h / 2 | |
| def _iou(a: list[float], b: list[float]) -> float: | |
| ax1, ay1, ax2, ay2 = _xywhn_to_corners(a) | |
| bx1, by1, bx2, by2 = _xywhn_to_corners(b) | |
| ix1, iy1 = max(ax1, bx1), max(ay1, by1) | |
| ix2, iy2 = min(ax2, bx2), min(ay2, by2) | |
| iw, ih = max(0.0, ix2 - ix1), max(0.0, iy2 - iy1) | |
| inter = iw * ih | |
| if inter == 0: | |
| return 0.0 | |
| area_a = (ax2 - ax1) * (ay2 - ay1) | |
| area_b = (bx2 - bx1) * (by2 - by1) | |
| return inter / (area_a + area_b - inter) | |
| def _center_in(inner: list[float], outer: list[float]) -> bool: | |
| """True if the centre of `inner` lies inside the box `outer`.""" | |
| x1, y1, x2, y2 = _xywhn_to_corners(outer) | |
| return x1 <= inner[0] <= x2 and y1 <= inner[1] <= y2 | |
| # --------------------------------------------------------------------------- | |
| # Compliance settings | |
| # --------------------------------------------------------------------------- | |
| # Minimum IoU for a PPE box to count as "covering" a body-part box. | |
| MATCH_IOU = 0.10 | |
| # Used when the YAML does not set compliance_conf. | |
| DEFAULT_COMPLIANCE_CONF = 0.40 | |
| PERSON = "person" | |
| LEVELS = ("required", "advisory") | |
| # Body parts a person has two of. One protected hand says nothing about the other, | |
| # so PPE elsewhere on the same person cannot soften a bare one to "uncertain". | |
| PAIRED_PARTS = {"hands", "foot", "ear"} | |
| _RANK = {"present": 0, "uncertain": 1, "missing": 2} | |
| def load_ppe_rules(cfg: dict, class_names) -> list[PpeRule]: | |
| """Read ppe_rules (new format) or ppe_pairs (old format) from the YAML dict.""" | |
| rules: list[PpeRule] = [] | |
| raw_rules = cfg.get("ppe_rules") | |
| if raw_rules: | |
| for r in raw_rules: | |
| any_of = r["any_of"] | |
| any_of = [any_of] if isinstance(any_of, str) else list(any_of) | |
| level = str(r.get("level", "required")).lower() | |
| if level not in LEVELS: | |
| raise ValueError( | |
| f"ppe_rules: level must be one of {LEVELS}, got '{level}'" | |
| ) | |
| rules.append( | |
| PpeRule( | |
| item=str(r.get("item", any_of[0])), | |
| body_part=r["body_part"], | |
| any_of=any_of, | |
| level=level, | |
| ) | |
| ) | |
| else: | |
| # Old format: every listed item is its own required rule. | |
| for part, exp in (cfg.get("ppe_pairs") or {}).items(): | |
| for e in ([exp] if isinstance(exp, str) else list(exp)): | |
| rules.append(PpeRule(item=e, body_part=part, any_of=[e])) | |
| known = set(class_names) | |
| for r in rules: | |
| unknown = [c for c in (r.body_part, *r.any_of) if c not in known] | |
| if unknown: | |
| raise ValueError( | |
| f"PPE rule '{r.item}' uses class name(s) not in the YAML names: {unknown}" | |
| ) | |
| return rules | |
| class SafetyDetector: | |
| def __init__(self, weights: str, data_yaml: str, conf: float = 0.25, | |
| augment: bool = False, compliance_conf: float | None = None): | |
| self.model = YOLO(weights) | |
| self.conf = conf | |
| self.augment = augment # test-time augmentation (slower, higher recall) | |
| with open(data_yaml) as f: | |
| cfg = yaml.safe_load(f) | |
| # Class names come from the YAML, NOT from the names embedded in the .pt. | |
| # The weights were trained with a mis-ordered names list, so the embedded | |
| # names are wrong. The label indices are correct; only the names were off. | |
| # The YAML (official SH17 index order) is the single source of truth. | |
| raw_names = cfg["names"] | |
| if isinstance(raw_names, dict): | |
| self.names = {int(k): v for k, v in raw_names.items()} | |
| else: | |
| self.names = dict(enumerate(raw_names)) | |
| n_model = len(self.model.names) | |
| if len(self.names) != n_model: | |
| raise ValueError( | |
| f"{data_yaml} lists {len(self.names)} classes but the model has {n_model}" | |
| ) | |
| self.model.model.names = self.names # so annotated images use correct labels | |
| # What to check, and how sure a box must be to count as evidence. | |
| self.rules: list[PpeRule] = load_ppe_rules(cfg, self.names.values()) | |
| if compliance_conf is None: | |
| compliance_conf = cfg.get("compliance_conf", DEFAULT_COMPLIANCE_CONF) | |
| self.compliance_conf = float(compliance_conf) | |
| # Kept for older code that reads detector.ppe_pairs. | |
| self.ppe_pairs: dict[str, list[str]] = {} | |
| for r in self.rules: | |
| self.ppe_pairs.setdefault(r.body_part, []).extend(r.any_of) | |
| def detect(self, image_path: str) -> DetectionResult: | |
| _, det_result = self._run(image_path) | |
| return det_result | |
| def detect_and_annotate(self, image, name: str | None = None): | |
| """ | |
| Same as detect(), but also returns a BGR numpy array with the detection | |
| boxes drawn on it. Runs YOLO only once. | |
| `image` can be a file path OR an in-memory BGR array. | |
| """ | |
| result, det_result = self._run(image, name=name) | |
| annotated = result.plot() | |
| return det_result, annotated | |
| def _run(self, image, name: str | None = None): | |
| result = self.model(image, conf=self.conf, augment=self.augment, | |
| verbose=False)[0] | |
| result.names = self.names # guarantee correct labels in plot + JSON | |
| names = self.names | |
| dets: list[Detection] = [] | |
| for box in result.boxes: | |
| cls_name = names[int(box.cls)] | |
| dets.append( | |
| Detection( | |
| cls=cls_name, | |
| confidence=round(float(box.conf), 3), | |
| box_xywhn=[round(float(x), 4) for x in box.xywhn[0].tolist()], | |
| ) | |
| ) | |
| compliance = self._assess_compliance(dets) | |
| people = len(self._people(dets)) | |
| if name is None: | |
| name = str(Path(image).name) if isinstance(image, (str, Path)) else "image" | |
| det_result = DetectionResult(image=name, detections=dets, | |
| compliance=compliance, people=people) | |
| return result, det_result | |
| # ----------------------------------------------------------------------- | |
| # Compliance | |
| # ----------------------------------------------------------------------- | |
| def _people(self, dets: list[Detection]) -> list[Detection]: | |
| """Confidently detected people, numbered left to right (index 0 = person 1).""" | |
| people = [d for d in dets | |
| if d.cls == PERSON and d.confidence >= self.compliance_conf] | |
| return sorted(people, key=lambda d: d.box_xywhn[0]) | |
| def _assess_compliance(self, dets: list[Detection]) -> list[ComplianceFinding]: | |
| """ | |
| For each person, and each PPE rule, decide present / missing / uncertain | |
| from the boxes that belong to THAT person. Purely rule-based and explainable. | |
| """ | |
| people = self._people(dets) | |
| # Give every non-person box to the person whose box contains its centre | |
| # (the smallest such box if people overlap). None = belongs to nobody. | |
| groups: dict[int | None, list[Detection]] = {} | |
| for d in dets: | |
| if d.cls == PERSON: | |
| continue | |
| inside = [i for i, p in enumerate(people) | |
| if _center_in(d.box_xywhn, p.box_xywhn)] | |
| if inside: | |
| i = min(inside, | |
| key=lambda k: people[k].box_xywhn[2] * people[k].box_xywhn[3]) | |
| pid: int | None = i + 1 | |
| else: | |
| pid = None | |
| groups.setdefault(pid, []).append(d) | |
| findings: list[ComplianceFinding] = [] | |
| for pid in [*range(1, len(people) + 1), None]: | |
| items = groups.get(pid, []) | |
| for rule in self.rules: | |
| ppe = [d for d in items if d.cls in rule.any_of] | |
| if rule.body_part == PERSON: | |
| if pid is None: | |
| continue # a vest that belongs to nobody: nothing to assess | |
| parts = [people[pid - 1]] | |
| else: | |
| parts = [d for d in items if d.cls == rule.body_part] | |
| outcome = self._judge(rule, parts, ppe, scoped=pid is not None) | |
| if outcome is None: | |
| continue # neither body part nor PPE seen: not assessed | |
| status, detail, conf = outcome | |
| findings.append( | |
| ComplianceFinding( | |
| body_part=rule.body_part, | |
| expected_ppe=" or ".join(rule.any_of), | |
| status=status, | |
| detail=detail, | |
| item=rule.item, | |
| level=rule.level, | |
| required=rule.level == "required", | |
| person_id=pid, | |
| confidence=conf, | |
| ) | |
| ) | |
| return findings | |
| def _judge(self, rule: PpeRule, parts: list[Detection], ppe: list[Detection], | |
| scoped: bool): | |
| """ | |
| Decide one rule for one person. | |
| `parts` body-part boxes of this person (for the vest rule: the person box) | |
| `ppe` PPE boxes of this person that satisfy the rule | |
| `scoped` False when the boxes could not be tied to a single person | |
| Returns (status, detail, confidence) or None when there is nothing to assess. | |
| """ | |
| if not parts and not ppe: | |
| return None | |
| if not parts and not scoped: | |
| # A PPE box that belongs to nobody (gloves on a bench, a stray | |
| # low-confidence box) says nothing about anyone's compliance. | |
| return None | |
| floor = self.compliance_conf | |
| label = " or ".join(rule.any_of) | |
| part_name = rule.body_part | |
| on_person = part_name == PERSON # vest: being on the person IS the match | |
| def match(part: Detection, p: Detection) -> float: | |
| return 1.0 if on_person else _iou(part.box_xywhn, p.box_xywhn) | |
| strong = [p for p in ppe if p.confidence >= floor] | |
| weak = [p for p in ppe if p.confidence < floor] | |
| where = "on person" if on_person else f"over {part_name}" | |
| outcomes: list[tuple[str, str, float]] = [] | |
| for part in parts: | |
| s_iou, s_box = max(((match(part, p), p) for p in strong), | |
| key=lambda t: t[0], default=(0.0, None)) | |
| w_iou, w_box = max(((match(part, p), p) for p in weak), | |
| key=lambda t: t[0], default=(0.0, None)) | |
| if s_box is not None and s_iou >= MATCH_IOU: | |
| outcomes.append(( | |
| "present", | |
| f"{s_box.cls} ({s_box.confidence:.2f}) {where}", | |
| s_box.confidence, | |
| )) | |
| elif w_box is not None and w_iou >= MATCH_IOU: | |
| outcomes.append(( | |
| "uncertain", | |
| f"low-confidence {w_box.cls} ({w_box.confidence:.2f}) {where}", | |
| w_box.confidence, | |
| )) | |
| elif (strong and scoped and len(parts) == 1 | |
| and part_name not in PAIRED_PARTS): | |
| outcomes.append(( | |
| "uncertain", | |
| f"{label} detected on this person but not clearly over the " | |
| f"{part_name} (max IoU {s_iou:.2f})", | |
| max(p.confidence for p in strong), | |
| )) | |
| elif part.confidence < floor: | |
| outcomes.append(( | |
| "uncertain", | |
| f"low-confidence {part_name} detection ({part.confidence:.2f}), " | |
| f"no {label} found", | |
| part.confidence, | |
| )) | |
| else: | |
| outcomes.append(( | |
| "missing", | |
| f"{part_name} detected ({part.confidence:.2f}) with no {label}", | |
| part.confidence, | |
| )) | |
| # PPE on the person that sits on no detected body part. A gloved hand is | |
| # usually not detected as "hands" at all, so the glove box is the evidence. | |
| unmatched = [p for p in ppe | |
| if not any(match(part, p) >= MATCH_IOU for part in parts)] | |
| for p in unmatched: | |
| if p.confidence >= floor: | |
| outcomes.append(( | |
| "present", f"{p.cls} detected ({p.confidence:.2f})", p.confidence | |
| )) | |
| if not outcomes: | |
| # Only low-confidence PPE boxes and no body part. | |
| p = max(weak, key=lambda d: d.confidence) | |
| outcomes.append(( | |
| "uncertain", | |
| f"low-confidence {p.cls} detection ({p.confidence:.2f})", | |
| p.confidence, | |
| )) | |
| # Worst case wins: one bare hand outweighs one gloved hand. | |
| worst = max(outcomes, key=lambda o: _RANK[o[0]]) | |
| details: list[str] = [] | |
| for _, text, _ in sorted(outcomes, key=lambda o: -_RANK[o[0]]): | |
| if text not in details: | |
| details.append(text) | |
| return worst[0], "; ".join(details), round(float(worst[2]), 3) |