Spaces:
Paused
Paused
File size: 5,684 Bytes
5a1cbfa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 | """
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
@dataclass
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)
@dataclass
class ComplianceFinding:
body_part: str
expected_ppe: str
status: str # "present" | "missing" | "uncertain"
detail: str
def to_dict(self) -> dict:
return asdict(self)
@dataclass
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
|