| """PPE detector miner (SN44 public track, manak0/Detect-worksite-ppe-compliance-detection-probe-set). |
| |
| Ensemble: every YOLO graph in weights*.onnx (FP32 or INT8 QDQ; one merged file may hold several, input i -> output i) runs on the frame, candidates (conf >= 0.05, |
| per-model NMS 0.7) are fused with weighted box fusion, then per-class thresholds + NMS are applied. The CPU provider keeps inference consistent between the local |
| compliance check and the chute output (IoU 0.95). |
| No file/network access here (sandbox rules): thresholds below are written by scripts/ppe/package_ens.py. |
| """ |
| from pathlib import Path |
|
|
| import cv2 |
| import numpy as np |
| import onnxruntime as ort |
| from numpy import ndarray |
| from pydantic import BaseModel |
|
|
| MANIFEST_CLASSES = ["safety helmet", "high-visibility safety vest", "protective gloves", "protective eyewear"] |
| CONF_THRESHOLD = 0.3 |
| PER_CLASS_CONF: dict[str, float] = {"safety helmet": 0.3, "high-visibility safety vest": 0.3, "protective gloves": 0.35, "protective eyewear": 0.275} |
| NMS_IOU_THRESHOLD = 0.45 |
| EYE_ALONE_CONF = 0.5 |
| SCENE_EYE_CONF = 0.7 |
| MAX_DETECTIONS = 100 |
| CAND_CONF, CAND_NMS, CAND_MAX, WBF_IOU = 0.05, 0.7, 300, 0.55 |
|
|
|
|
| class BoundingBox(BaseModel): |
| x1: int |
| y1: int |
| x2: int |
| y2: int |
| cls_id: int |
| conf: float |
|
|
|
|
| class TVFrameResult(BaseModel): |
| frame_id: int |
| boxes: list[BoundingBox] | None = None |
| polygons: list | None = None |
| keypoints: list[tuple[int, int]] | None = None |
|
|
|
|
| def _letterbox(img: ndarray, size: int): |
| h, w = img.shape[:2] |
| r = min(size / h, size / w) |
| nw, nh = int(round(w * r)), int(round(h * r)) |
| px, py = (size - nw) // 2, (size - nh) // 2 |
| canvas = np.full((size, size, 3), 114, dtype=np.uint8) |
| canvas[py:py + nh, px:px + nw] = cv2.resize(img, (nw, nh), interpolation=cv2.INTER_LINEAR) |
| return canvas, r, px, py |
|
|
|
|
| def _nms(b: ndarray, s: ndarray, thr: float) -> list[int]: |
| area = (b[:, 2] - b[:, 0]) * (b[:, 3] - b[:, 1]) |
| order, keep = s.argsort()[::-1], [] |
| while order.size: |
| i = int(order[0]); keep.append(i); rest = order[1:] |
| iw = np.clip(np.minimum(b[i, 2], b[rest, 2]) - np.maximum(b[i, 0], b[rest, 0]), 0, None) |
| ih = np.clip(np.minimum(b[i, 3], b[rest, 3]) - np.maximum(b[i, 1], b[rest, 1]), 0, None) |
| inter = iw * ih |
| order = rest[inter / np.maximum(area[i] + area[rest] - inter, 1e-9) <= thr] |
| return keep |
|
|
|
|
| def _wbf(lists: list[list[list[float]]], thr: float) -> list[list[float]]: |
| """lists: one list of [x1, y1, x2, y2, cls, conf] per model; conf = summed conf / number of models.""" |
| allb = sorted((b for lst in lists for b in lst), key=lambda b: -b[5]) |
| clusters: list[list] = [] |
| for b in allb: |
| best, bv = None, thr |
| for cl in clusters: |
| f = cl[0] |
| if f[4] == b[4]: |
| iw = max(0.0, min(f[2], b[2]) - max(f[0], b[0])); ih = max(0.0, min(f[3], b[3]) - max(f[1], b[1])) |
| inter = iw * ih; u = (f[2] - f[0]) * (f[3] - f[1]) + (b[2] - b[0]) * (b[3] - b[1]) - inter |
| v = inter / u if u > 0 else 0.0 |
| if v >= bv: |
| best, bv = cl, v |
| if best is None: |
| clusters.append([list(b), [b]]) |
| else: |
| best[1].append(b) |
| ws = np.array([m[5] for m in best[1]]); co = np.array([m[:4] for m in best[1]], float) |
| best[0][:4] = list((co * ws[:, None]).sum(0) / ws.sum()) |
| return [[*cl[0][:4], cl[0][4], min(1.0, sum(m[5] for m in cl[1]) / len(lists))] for cl in clusters] |
|
|
|
|
| def _alone_eyewear(e: ndarray, helmets: list[ndarray], vests: list[ndarray]) -> bool: |
| """No helmet above this eyewear box and no vest below it: ordinary glasses and sunglasses on people without PPE are mostly not labelled.""" |
| x, y = (e[0] + e[2]) / 2, (e[1] + e[3]) / 2 |
| for b in helmets: |
| w, h = b[2] - b[0], b[3] - b[1] |
| if b[0] - .3 * w <= x <= b[2] + .3 * w and b[1] - 5 <= y <= b[3] + 1.2 * h: |
| return False |
| for b in vests: |
| w, h = b[2] - b[0], b[3] - b[1] |
| if b[0] - .25 * w <= x <= b[2] + .25 * w and y <= b[1] + .3 * h and b[1] - y <= 1.2 * w: |
| return False |
| return True |
|
|
|
|
| class Miner: |
| def __init__(self, path_hf_repo: Path) -> None: |
| so = ort.SessionOptions() |
| so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL |
| so.intra_op_num_threads = 2 |
| so.inter_op_num_threads = 1 |
| so.add_session_config_entry("session.set_denormal_as_zero", "1") |
| so.add_session_config_entry("session.intra_op.allow_spinning", "0") |
| self.sessions = [ort.InferenceSession(str(p), so, providers=["CPUExecutionProvider"]) |
| for p in sorted(Path(path_hf_repo).glob("weights*.onnx"))] |
| |
| self.inputs = [(k, i.name, int(i.shape[2]), s.get_outputs()[j].name) |
| for k, s in enumerate(self.sessions) for j, i in enumerate(s.get_inputs())] |
| self.cls_conf = np.array([PER_CLASS_CONF.get(n, CONF_THRESHOLD) for n in MANIFEST_CLASSES], np.float32) |
| for _ in range(2): |
| self._infer(np.zeros((640, 640, 3), np.uint8)) |
|
|
| def __repr__(self) -> str: |
| return f"PPE YOLO ONNX ensemble inputs={[b[2] for b in self.inputs]}" |
|
|
| def _candidates(self, img: ndarray) -> list[list[list[float]]]: |
| feeds: list[dict] = [{} for _ in self.sessions] |
| geo = [] |
| for k, name, size, _ in self.inputs: |
| canvas, r, px, py = _letterbox(img, size) |
| feeds[k][name] = np.ascontiguousarray(canvas[:, :, ::-1].transpose(2, 0, 1)[None], dtype=np.float32) / 255.0 |
| geo.append((r, px, py)) |
| outs = {} |
| for k, s in enumerate(self.sessions): |
| names = [b[3] for b in self.inputs if b[0] == k] |
| outs.update(zip(((k, n) for n in names), s.run(names, feeds[k]))) |
| return [self._decode(outs[(k, oname)][0], img, *geo[b]) for b, (k, _, _, oname) in enumerate(self.inputs)] |
|
|
| def _decode(self, out: ndarray, img: ndarray, r: float, px: int, py: int) -> list[list[float]]: |
| if out.shape[-1] == 6: |
| boxes, scores, cls = out[:, :4], out[:, 4], out[:, 5].astype(np.int32) |
| else: |
| out = out.T |
| c = out[:, 4:] |
| cls = c.argmax(1).astype(np.int32) |
| scores = c[np.arange(len(c)), cls] |
| xy, wh = out[:, :2], out[:, 2:4] / 2 |
| boxes = np.concatenate([xy - wh, xy + wh], 1) |
| ok = (cls >= 0) & (cls < len(MANIFEST_CLASSES)) & (scores >= CAND_CONF) |
| boxes, scores, cls = boxes[ok].astype(np.float32), scores[ok], cls[ok] |
| boxes[:, [0, 2]] = ((boxes[:, [0, 2]] - px) / r).clip(0, img.shape[1]) |
| boxes[:, [1, 3]] = ((boxes[:, [1, 3]] - py) / r).clip(0, img.shape[0]) |
| res = [] |
| for c_ in np.unique(cls): |
| sel = np.where(cls == c_)[0] |
| res += [[*map(float, boxes[sel[j]]), int(c_), float(scores[sel[j]])] |
| for j in _nms(boxes[sel], scores[sel], CAND_NMS)] |
| res.sort(key=lambda b: -b[5]) |
| return res[:CAND_MAX] |
|
|
| def _infer(self, img: ndarray) -> list[BoundingBox]: |
| cands = self._candidates(img) |
| |
| fused = cands[0] if len(cands) == 1 else _wbf(cands, WBF_IOU) |
| fused = [b for b in fused if b[5] >= self.cls_conf[b[4]]] |
| if not fused: |
| return [] |
| boxes = np.array([b[:4] for b in fused], np.float32); scores = np.array([b[5] for b in fused], np.float32) |
| cls = np.array([b[4] for b in fused], np.int32) |
| kept = [] |
| for k in np.unique(cls): |
| sel = np.where(cls == k)[0] |
| kept += [int(sel[j]) for j in _nms(boxes[sel], scores[sel], NMS_IOU_THRESHOLD)] |
| helmets = [boxes[i] for i in kept if cls[i] == 0]; vests = [boxes[i] for i in kept if cls[i] == 1] |
| res = [] |
| for i in kept: |
| x1, y1, x2, y2 = boxes[i] |
| if cls[i] == 3 and scores[i] < EYE_ALONE_CONF and _alone_eyewear(boxes[i], helmets, vests): |
| continue |
| if cls[i] == 3 and scores[i] < SCENE_EYE_CONF and not helmets and not vests: |
| continue |
| if x2 - x1 >= 1 and y2 - y1 >= 1: |
| res.append(BoundingBox(x1=int(round(float(x1))), y1=int(round(float(y1))), x2=int(round(float(x2))), |
| y2=int(round(float(y2))), cls_id=int(cls[i]), conf=float(min(1.0, scores[i])))) |
| res.sort(key=lambda b: -b.conf) |
| return res[:MAX_DETECTIONS] |
|
|
| def predict_batch(self, batch_images: list[ndarray], offset: int, n_keypoints: int) -> list[TVFrameResult]: |
| kp = [(0, 0)] * max(0, int(n_keypoints)) |
| out = [] |
| for i, img in enumerate(batch_images): |
| try: |
| boxes = self._infer(img) |
| except Exception as exc: |
| print(f"inference failed on frame {offset + i}: {exc}") |
| boxes = [] |
| out.append(TVFrameResult(frame_id=offset + i, boxes=boxes, polygons=[], keypoints=list(kp))) |
| return out |
|
|
