| """TurboVision public-track miner for manak0/Detect-worksite-ppe-compliance-detection-probe-set. |
| |
| Model: YOLO26n exported to ONNX (input 1x3x704x704, NMS embedded, output 1x300x6 |
| rows of x1, y1, x2, y2, score, class in letterboxed pixel space). |
| |
| cls_id in the response is the index into the element's manifest `objects` list. |
| """ |
|
|
| 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", |
| ] |
|
|
| |
| CLASS_CONF = np.array([0.40, 0.40, 0.40, 0.30], dtype=np.float32) |
|
|
| MODEL_FILENAME = "weights.onnx" |
| LETTERBOX_FILL = 114 |
|
|
|
|
| class BoundingBox(BaseModel): |
| x1: int |
| y1: int |
| x2: int |
| y2: int |
| cls_id: int |
| conf: float |
|
|
|
|
| class Polygon(BaseModel): |
| cls_id: int |
| conf: float |
| points: list[tuple[int, int]] |
|
|
|
|
| class TVFrameResult(BaseModel): |
| frame_id: int |
| boxes: list[BoundingBox] | None = None |
| polygons: list[Polygon] | None = None |
| keypoints: list[tuple[int, int]] | None = None |
|
|
|
|
| def _letterbox(image: ndarray, size: int) -> tuple[ndarray, float, int, int]: |
| """Ultralytics-style centred letterbox to a size x size canvas.""" |
| height, width = image.shape[:2] |
| ratio = min(size / height, size / width) |
| new_w = int(round(width * ratio)) |
| new_h = int(round(height * ratio)) |
| if (new_w, new_h) != (width, height): |
| image = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LINEAR) |
| pad_w = (size - new_w) / 2 |
| pad_h = (size - new_h) / 2 |
| top, bottom = int(round(pad_h - 0.1)), int(round(pad_h + 0.1)) |
| left, right = int(round(pad_w - 0.1)), int(round(pad_w + 0.1)) |
| canvas = cv2.copyMakeBorder( |
| image, |
| top, |
| bottom, |
| left, |
| right, |
| cv2.BORDER_CONSTANT, |
| value=(LETTERBOX_FILL, LETTERBOX_FILL, LETTERBOX_FILL), |
| ) |
| return canvas, ratio, left, top |
|
|
|
|
| class Miner: |
| def __init__(self, path_hf_repo: Path) -> None: |
| model_path = Path(path_hf_repo) / MODEL_FILENAME |
| if not model_path.is_file(): |
| raise FileNotFoundError(f"{MODEL_FILENAME} not found in {path_hf_repo}") |
|
|
| options = ort.SessionOptions() |
| options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL |
| options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL |
| options.intra_op_num_threads = 2 |
| options.inter_op_num_threads = 1 |
|
|
| providers = ["CPUExecutionProvider"] |
| if "CUDAExecutionProvider" in ort.get_available_providers(): |
| providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] |
| try: |
| self.session = ort.InferenceSession( |
| str(model_path), sess_options=options, providers=providers |
| ) |
| except Exception as exc: |
| print(f"CUDA session failed, falling back to CPU: {exc}") |
| self.session = ort.InferenceSession( |
| str(model_path), sess_options=options, providers=["CPUExecutionProvider"] |
| ) |
|
|
| model_input = self.session.get_inputs()[0] |
| self.input_name = model_input.name |
| self.input_size = int(model_input.shape[2]) |
| self.output_name = self.session.get_outputs()[0].name |
| self._check_class_order() |
|
|
| warmup = np.zeros((self.input_size, self.input_size, 3), dtype=np.uint8) |
| for _ in range(2): |
| self._infer_one(warmup) |
| print(f"Miner ready: providers={self.session.get_providers()} input={self.input_size}") |
|
|
| def __repr__(self) -> str: |
| return ( |
| f"PPE YOLO26n ONNX input={self.input_size} " |
| f"providers={self.session.get_providers()}" |
| ) |
|
|
| def _check_class_order(self) -> None: |
| names = self.session.get_modelmeta().custom_metadata_map.get("names", "") |
| for idx, name in enumerate(MANIFEST_CLASSES): |
| if f"{idx}: '{name}'" not in names: |
| raise ValueError(f"ONNX class order does not match manifest: {names}") |
|
|
| def _infer_one(self, image: ndarray) -> list[BoundingBox]: |
| canvas, ratio, pad_x, pad_y = _letterbox(image, self.input_size) |
| rgb = cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB) |
| tensor = np.ascontiguousarray(rgb.transpose(2, 0, 1)[None], dtype=np.float32) / 255.0 |
| raw = self.session.run([self.output_name], {self.input_name: tensor})[0][0] |
|
|
| scores = raw[:, 4] |
| cls_ids = raw[:, 5].astype(np.int64) |
| known = (cls_ids >= 0) & (cls_ids < len(MANIFEST_CLASSES)) |
| keep = known & (scores >= CLASS_CONF[np.clip(cls_ids, 0, len(CLASS_CONF) - 1)]) |
| if not keep.any(): |
| return [] |
|
|
| boxes = raw[keep, :4].astype(np.float64) |
| boxes[:, [0, 2]] = (boxes[:, [0, 2]] - pad_x) / ratio |
| boxes[:, [1, 3]] = (boxes[:, [1, 3]] - pad_y) / ratio |
| height, width = image.shape[:2] |
| boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, width) |
| boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, height) |
|
|
| results: list[BoundingBox] = [] |
| for (x1, y1, x2, y2), score, cls_id in zip(boxes, scores[keep], cls_ids[keep]): |
| x1, y1, x2, y2 = (int(round(v)) for v in (x1, y1, x2, y2)) |
| if x2 <= x1 or y2 <= y1: |
| continue |
| results.append( |
| BoundingBox( |
| x1=x1, y1=y1, x2=x2, y2=y2, cls_id=int(cls_id), conf=float(score) |
| ) |
| ) |
| return results |
|
|
| def predict_batch( |
| self, |
| batch_images: list[ndarray], |
| offset: int, |
| n_keypoints: int, |
| ) -> list[TVFrameResult]: |
| results: list[TVFrameResult] = [] |
| for index, image in enumerate(batch_images): |
| try: |
| boxes = self._infer_one(image) |
| except Exception as exc: |
| print(f"Inference failed on frame {offset + index}: {exc}") |
| boxes = [] |
| results.append( |
| TVFrameResult( |
| frame_id=offset + index, |
| boxes=boxes, |
| polygons=[], |
| keypoints=[(0, 0)] * max(0, int(n_keypoints)), |
| ) |
| ) |
| return results |
|
|
