Spaces:
Paused
Paused
| """ | |
| 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: | |
| "a head is detected but no helmet overlaps it" -> possible missing helmet | |
| The VLM never decides compliance on its own — it is handed these facts. This keeps | |
| the safety-critical judgement in code that can be logged and reviewed, in line with | |
| the AiGiS design principle. | |
| """ | |
| 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 | |
| status: str # "present" | "missing" | "uncertain" | |
| detail: str | |
| 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) | |
| def to_dict(self) -> dict: | |
| return { | |
| "image": self.image, | |
| "detections": [d.to_dict() for d in self.detections], | |
| "compliance": [c.to_dict() for c in self.compliance], | |
| } | |
| 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) | |
| class SafetyDetector: | |
| def __init__(self, weights: str, data_yaml: str, conf: float = 0.25): | |
| self.model = YOLO(weights) | |
| self.conf = conf | |
| with open(data_yaml) as f: | |
| cfg = yaml.safe_load(f) | |
| # body_part -> expected PPE class name | |
| self.ppe_pairs: dict[str, str] = cfg.get("ppe_pairs", {}) | |
| 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 (via Ultralytics' own .plot()). Runs YOLO only once, | |
| so callers can get both the JSON and the annotated image from one call. | |
| `image` can be a file path OR an in-memory BGR array (e.g. from | |
| cv2.imdecode on an uploaded file) — Ultralytics accepts both. | |
| """ | |
| 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, verbose=False)[0] | |
| names = result.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) | |
| 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) | |
| return result, det_result | |
| def _assess_compliance(self, dets: list[Detection]) -> list[ComplianceFinding]: | |
| """ | |
| For each body part that should be protected, check whether the expected | |
| PPE was detected overlapping it. Purely rule-based and explainable. | |
| """ | |
| findings: list[ComplianceFinding] = [] | |
| by_cls: dict[str, list[Detection]] = {} | |
| for d in dets: | |
| by_cls.setdefault(d.cls, []).append(d) | |
| for body_part, expected in self.ppe_pairs.items(): | |
| parts = by_cls.get(body_part, []) | |
| if not parts: | |
| continue # body part not visible -> nothing to assess | |
| ppe_items = by_cls.get(expected, []) | |
| for part in parts: | |
| overlap = max( | |
| (_iou(part.box_xywhn, p.box_xywhn) for p in ppe_items), | |
| default=0.0, | |
| ) | |
| if overlap >= 0.10: | |
| status, detail = "present", ( | |
| f"{expected} overlaps detected {body_part} (IoU {overlap:.2f})" | |
| ) | |
| elif ppe_items: | |
| status, detail = "uncertain", ( | |
| f"{expected} present in image but not clearly over this " | |
| f"{body_part} (max IoU {overlap:.2f})" | |
| ) | |
| else: | |
| status, detail = "missing", ( | |
| f"{body_part} detected but no {expected} found in image" | |
| ) | |
| findings.append( | |
| ComplianceFinding(body_part, expected, status, detail) | |
| ) | |
| return findings | |