aigis-captioning / detector.py
ZahraFatima's picture
Update detector.py
8c57ad7 verified
Raw History Blame Contribute Delete
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
@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 # 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)
@dataclass
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],
}
@dataclass
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)