| from pathlib import Path |
| import argparse |
| import json |
| import math, cv2, numpy as np, onnxruntime as ort |
| from numpy import ndarray |
| from pydantic import BaseModel |
|
|
|
|
| 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] |
| keypoints: list[tuple[int, int]] |
|
|
|
|
| class Miner: |
| def __init__(self, path_hf_repo: Path) -> None: |
| model_path = path_hf_repo / "weights.onnx" |
|
|
| try: |
| ort.preload_dlls() |
| except Exception: |
| pass |
|
|
| sess_options = ort.SessionOptions() |
| sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL |
|
|
| try: |
| self.session = ort.InferenceSession( |
| str(model_path), |
| sess_options=sess_options, |
| providers=["CUDAExecutionProvider", "CPUExecutionProvider"], |
| ) |
| except Exception: |
| self.session = ort.InferenceSession( |
| str(model_path), |
| sess_options=sess_options, |
| providers=["CPUExecutionProvider"], |
| ) |
|
|
| self.input_name = self.session.get_inputs()[0].name |
| self.output_names = [o.name for o in self.session.get_outputs()] |
| self.input_shape = self.session.get_inputs()[0].shape |
|
|
| self.input_height = self._safe_dim(self.input_shape[2], 1280) |
| self.input_width = self._safe_dim(self.input_shape[3], 1280) |
|
|
| self.tile_size = 960 |
| self.overlap = 0.5 |
|
|
| self.conf_thres = 0.47 |
| self.wbf_iou = 0.55 |
| self.max_det = 100 |
|
|
| self.min_box_area = 14 * 14 |
| self.min_w = 8 |
| self.min_h = 8 |
| self.max_aspect_ratio = 8.0 |
| self.max_box_area_ratio = 0.8 |
|
|
| def __repr__(self) -> str: |
| return ( |
| f"ONNXRuntime(session={type(self.session).__name__}, " |
| f"providers={self.session.get_providers()})" |
| ) |
|
|
| @staticmethod |
| def _safe_dim(value, default: int) -> int: |
| return value if isinstance(value, int) and value > 0 else default |
|
|
| @staticmethod |
| def _clip_boxes(boxes: np.ndarray, image_size: tuple[int, int]) -> np.ndarray: |
| w, h = image_size |
| if len(boxes) == 0: |
| return boxes |
| boxes[:, 0] = np.clip(boxes[:, 0], 0, w - 1) |
| boxes[:, 1] = np.clip(boxes[:, 1], 0, h - 1) |
| boxes[:, 2] = np.clip(boxes[:, 2], 0, w - 1) |
| boxes[:, 3] = np.clip(boxes[:, 3], 0, h - 1) |
| return boxes |
|
|
| @staticmethod |
| def _xywh_to_xyxy(boxes: np.ndarray) -> np.ndarray: |
| out = np.empty_like(boxes) |
| out[:, 0] = boxes[:, 0] - boxes[:, 2] / 2.0 |
| out[:, 1] = boxes[:, 1] - boxes[:, 3] / 2.0 |
| out[:, 2] = boxes[:, 0] + boxes[:, 2] / 2.0 |
| out[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.0 |
| return out |
|
|
| def _slice_image( |
| self, image: np.ndarray |
| ) -> list[tuple[np.ndarray, tuple[int, int], tuple[int, int]]]: |
| h, w = image.shape[:2] |
| t = self.tile_size |
| st = max(1, int(t * (1.0 - self.overlap))) |
|
|
| xs = [] |
| x = 0 |
| while True: |
| if x + t >= w: |
| xs.append(max(0, w - t)) |
| break |
| xs.append(x) |
| x += st |
|
|
| ys = [] |
| y = 0 |
| while True: |
| if y + t >= h: |
| ys.append(max(0, h - t)) |
| break |
| ys.append(y) |
| y += st |
|
|
| xs = list(dict.fromkeys(xs)) |
| ys = list(dict.fromkeys(ys)) |
|
|
| out = [] |
| for y0 in ys: |
| for x0 in xs: |
| x1 = min(x0 + t, w) |
| y1 = min(y0 + t, h) |
| crop = image[y0:y1, x0:x1] |
| vh, vw = crop.shape[:2] |
| out.append((crop, (x0, y0), (vw, vh))) |
| return out |
|
|
| def _preprocess(self, image: np.ndarray) -> np.ndarray: |
| if image.dtype != np.uint8: |
| image = image.astype(np.uint8) |
|
|
| h, w = image.shape[:2] |
| pad_w = max(0, self.input_width - w) |
| pad_h = max(0, self.input_height - h) |
|
|
| if pad_w > 0 or pad_h > 0: |
| image = cv2.copyMakeBorder( |
| image, |
| 0, |
| pad_h, |
| 0, |
| pad_w, |
| cv2.BORDER_CONSTANT, |
| value=(114, 114, 114), |
| ) |
|
|
| if image.shape[0] != self.input_height or image.shape[1] != self.input_width: |
| raise ValueError( |
| f"Bad tile shape after pad: {image.shape[:2]}, " |
| f"expected ({self.input_height}, {self.input_width})" |
| ) |
|
|
| img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) |
| img = img.astype(np.float32) / 255.0 |
| img = np.transpose(img, (2, 0, 1))[None, ...] |
| return np.ascontiguousarray(img, dtype=np.float32) |
|
|
| def _filter_sane_boxes( |
| self, |
| boxes: np.ndarray, |
| scores: np.ndarray, |
| cls_ids: np.ndarray, |
| valid_size: tuple[int, int], |
| ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: |
| if len(boxes) == 0: |
| return boxes, scores, cls_ids |
|
|
| valid_w, valid_h = valid_size |
| image_area = float(max(1, valid_w * valid_h)) |
|
|
| keep = [] |
| for i, box in enumerate(boxes): |
| x1, y1, x2, y2 = box.tolist() |
| bw = x2 - x1 |
| bh = y2 - y1 |
|
|
| if bw <= 0 or bh <= 0: |
| continue |
| if bw < self.min_w or bh < self.min_h: |
| continue |
|
|
| area = bw * bh |
| if area < self.min_box_area: |
| continue |
| if area > self.max_box_area_ratio * image_area: |
| continue |
|
|
| ar = max(bw / max(bh, 1e-6), bh / max(bw, 1e-6)) |
| if ar > self.max_aspect_ratio: |
| continue |
|
|
| keep.append(i) |
|
|
| if not keep: |
| return ( |
| np.empty((0, 4), dtype=np.float32), |
| np.empty((0,), dtype=np.float32), |
| np.empty((0,), dtype=np.int32), |
| ) |
|
|
| keep = np.array(keep, dtype=np.intp) |
| return boxes[keep], scores[keep], cls_ids[keep] |
|
|
| def _decode_final_dets( |
| self, preds: np.ndarray, valid_size: tuple[int, int] |
| ) -> tuple[np.ndarray, np.ndarray]: |
| if preds.ndim == 3 and preds.shape[0] == 1: |
| preds = preds[0] |
|
|
| if preds.ndim != 2 or preds.shape[1] < 6: |
| raise ValueError(f"Unexpected final-det output shape: {preds.shape}") |
|
|
| boxes = preds[:, :4].astype(np.float32) |
| scores = preds[:, 4].astype(np.float32) |
| cls_ids = preds[:, 5].astype(np.int32) |
|
|
| keep = cls_ids == 0 |
| boxes = boxes[keep] |
| scores = scores[keep] |
| cls_ids = cls_ids[keep] |
|
|
| keep = scores >= self.conf_thres |
| boxes = boxes[keep] |
| scores = scores[keep] |
| cls_ids = cls_ids[keep] |
|
|
| if len(boxes) == 0: |
| return ( |
| np.empty((0, 4), dtype=np.float32), |
| np.empty((0,), dtype=np.float32), |
| ) |
|
|
| boxes = self._clip_boxes(boxes, valid_size) |
| boxes, scores, cls_ids = self._filter_sane_boxes(boxes, scores, cls_ids, valid_size) |
|
|
| if len(boxes) == 0: |
| return ( |
| np.empty((0, 4), dtype=np.float32), |
| np.empty((0,), dtype=np.float32), |
| ) |
|
|
| return boxes, scores |
|
|
| def _decode_raw_yolo( |
| self, preds: np.ndarray, valid_size: tuple[int, int] |
| ) -> tuple[np.ndarray, np.ndarray]: |
| if preds.ndim != 3: |
| raise ValueError(f"Unexpected raw output shape: {preds.shape}") |
| if preds.shape[0] != 1: |
| raise ValueError(f"Unexpected batch dimension in raw output: {preds.shape}") |
|
|
| preds = preds[0] |
|
|
| if preds.shape[0] <= 16 and preds.shape[1] > preds.shape[0]: |
| preds = preds.T |
|
|
| if preds.ndim != 2 or preds.shape[1] < 5: |
| raise ValueError(f"Unexpected normalized raw output shape: {preds.shape}") |
|
|
| boxes_xywh = preds[:, :4].astype(np.float32) |
| tail = preds[:, 4:].astype(np.float32) |
|
|
| if tail.shape[1] == 1: |
| scores = tail[:, 0] |
| cls_ids = np.zeros(len(scores), dtype=np.int32) |
| elif tail.shape[1] == 2: |
| obj = tail[:, 0] |
| cls_prob = tail[:, 1] |
| scores = obj * cls_prob |
| cls_ids = np.zeros(len(scores), dtype=np.int32) |
| else: |
| obj = tail[:, 0] |
| class_probs = tail[:, 1:] |
| cls_ids = np.argmax(class_probs, axis=1).astype(np.int32) |
| cls_scores = class_probs[np.arange(len(class_probs)), cls_ids] |
| scores = obj * cls_scores |
|
|
| keep = cls_ids == 0 |
| boxes_xywh = boxes_xywh[keep] |
| scores = scores[keep] |
| cls_ids = cls_ids[keep] |
|
|
| keep = scores >= self.conf_thres |
| boxes_xywh = boxes_xywh[keep] |
| scores = scores[keep] |
| cls_ids = cls_ids[keep] |
|
|
| if len(boxes_xywh) == 0: |
| return ( |
| np.empty((0, 4), dtype=np.float32), |
| np.empty((0,), dtype=np.float32), |
| ) |
|
|
| boxes = self._xywh_to_xyxy(boxes_xywh) |
| boxes = self._clip_boxes(boxes, valid_size) |
| boxes, scores, cls_ids = self._filter_sane_boxes(boxes, scores, cls_ids, valid_size) |
|
|
| if len(boxes) == 0: |
| return ( |
| np.empty((0, 4), dtype=np.float32), |
| np.empty((0,), dtype=np.float32), |
| ) |
|
|
| return boxes, scores |
|
|
| def _postprocess( |
| self, output: np.ndarray, valid_size: tuple[int, int] |
| ) -> tuple[np.ndarray, np.ndarray]: |
| if output.ndim == 2 and output.shape[1] == 6: |
| return self._decode_final_dets(output, valid_size) |
|
|
| if output.ndim == 3 and output.shape[0] == 1 and output.shape[2] == 6: |
| return self._decode_final_dets(output, valid_size) |
|
|
| return self._decode_raw_yolo(output, valid_size) |
|
|
| def _predict_tile(self, image: np.ndarray) -> tuple[np.ndarray, np.ndarray]: |
| if image is None: |
| raise ValueError("Input image is None") |
| if not isinstance(image, np.ndarray): |
| raise TypeError(f"Input is not numpy array: {type(image)}") |
| if image.ndim != 3: |
| raise ValueError(f"Expected HWC image, got shape={image.shape}") |
| if image.shape[2] != 3: |
| raise ValueError(f"Expected 3 channels, got shape={image.shape}") |
|
|
| valid_w, valid_h = image.shape[1], image.shape[0] |
| input_tensor = self._preprocess(image) |
|
|
| expected_shape = (1, 3, self.input_height, self.input_width) |
| if input_tensor.shape != expected_shape: |
| raise ValueError( |
| f"Bad input tensor shape={input_tensor.shape}, expected={expected_shape}" |
| ) |
|
|
| outputs = self.session.run(self.output_names, {self.input_name: input_tensor}) |
| det_output = outputs[0] |
| return self._postprocess(det_output, (valid_w, valid_h)) |
|
|
| @staticmethod |
| def _iou_one(a: np.ndarray, b: np.ndarray) -> float: |
| xx1 = max(float(a[0]), float(b[0])) |
| yy1 = max(float(a[1]), float(b[1])) |
| xx2 = min(float(a[2]), float(b[2])) |
| yy2 = min(float(a[3]), float(b[3])) |
|
|
| inter = max(0.0, xx2 - xx1) * max(0.0, yy2 - yy1) |
| area_a = max(0.0, float(a[2] - a[0])) * max(0.0, float(a[3] - a[1])) |
| area_b = max(0.0, float(b[2] - b[0])) * max(0.0, float(b[3] - b[1])) |
| return float(inter / (area_a + area_b - inter + 1e-7)) |
|
|
| def _wbf(self, boxes: np.ndarray, scores: np.ndarray) -> tuple[np.ndarray, np.ndarray]: |
| if len(boxes) == 0: |
| return boxes, scores |
|
|
| order = np.argsort(scores)[::-1] |
| boxes = boxes[order] |
| scores = scores[order] |
|
|
| used = np.zeros(len(boxes), dtype=bool) |
| fused_boxes = [] |
| fused_scores = [] |
|
|
| for i in range(len(boxes)): |
| if used[i]: |
| continue |
|
|
| cluster = [i] |
| used[i] = True |
| changed = True |
|
|
| while changed: |
| changed = False |
| for j in range(i + 1, len(boxes)): |
| if used[j]: |
| continue |
| if any(self._iou_one(boxes[j], boxes[k]) >= self.wbf_iou for k in cluster): |
| used[j] = True |
| cluster.append(j) |
| changed = True |
|
|
| cb = boxes[cluster] |
| cs = scores[cluster] |
| w = cs / (cs.sum() + 1e-7) |
|
|
| fused_boxes.append((cb * w[:, None]).sum(axis=0)) |
| fused_scores.append(float(cs.max())) |
|
|
| return np.asarray(fused_boxes, dtype=np.float32), np.asarray(fused_scores, dtype=np.float32) |
|
|
| def _merge_boxes_global( |
| self, |
| image_shape: tuple[int, int], |
| tile_boxes: np.ndarray, |
| tile_scores: np.ndarray, |
| ) -> list[BoundingBox]: |
| if len(tile_boxes) == 0: |
| return [] |
|
|
| H, W = image_shape[:2] |
| boxes = tile_boxes.copy() |
| scores = tile_scores.copy() |
|
|
| keep = np.argsort(scores)[::-1] |
| boxes = boxes[keep] |
| scores = scores[keep] |
|
|
| boxes, scores = self._wbf(boxes, scores) |
|
|
| out = [] |
| for b, sc in zip(boxes[: self.max_det], scores[: self.max_det]): |
| x1 = int(math.floor(b[0])) |
| y1 = int(math.floor(b[1])) |
| x2 = int(math.ceil(b[2])) |
| y2 = int(math.ceil(b[3])) |
|
|
| x1 = max(0, min(x1, W - 1)) |
| y1 = max(0, min(y1, H - 1)) |
| x2 = max(0, min(x2, W - 1)) |
| y2 = max(0, min(y2, H - 1)) |
|
|
| if x2 > x1 and y2 > y1: |
| out.append( |
| BoundingBox( |
| x1=x1, |
| y1=y1, |
| x2=x2, |
| y2=y2, |
| cls_id=0, |
| conf=float(sc), |
| ) |
| ) |
| return out |
|
|
| def predict_image(self, image: np.ndarray) -> list[BoundingBox]: |
| if image is None: |
| raise ValueError("Input image is None") |
| if not isinstance(image, np.ndarray): |
| raise TypeError(f"Input is not numpy array: {type(image)}") |
| if image.ndim != 3: |
| raise ValueError(f"Expected HWC image, got shape={image.shape}") |
| if image.shape[2] != 3: |
| raise ValueError(f"Expected 3 channels, got shape={image.shape}") |
|
|
| H, W = image.shape[:2] |
| tiles = self._slice_image(image) |
|
|
| all_boxes = [] |
| all_scores = [] |
|
|
| for i, (tile_img, (ox, oy), (vw, vh)) in enumerate(tiles): |
| try: |
| b1, s1 = self._predict_tile(tile_img) |
|
|
| flipped = cv2.flip(tile_img, 1) |
| b2, s2 = self._predict_tile(flipped) |
|
|
| if len(b2): |
| tw = tile_img.shape[1] |
| b2 = b2.copy() |
| b2[:, [0, 2]] = tw - b2[:, [2, 0]] |
|
|
| if len(b1) == 0 and len(b2) == 0: |
| continue |
| if len(b1) == 0: |
| b, s = b2, s2 |
| elif len(b2) == 0: |
| b, s = b1, s1 |
| else: |
| b = np.vstack([b1, b2]) |
| s = np.hstack([s1, s2]) |
|
|
| except Exception as e: |
| print(f"⚠️ Tile inference failed at ({ox}, {oy}): {e}") |
| continue |
|
|
| left_edge = ox == 0 |
| top_edge = oy == 0 |
| right_edge = (ox + vw) >= W |
| bottom_edge = (oy + vh) >= H |
|
|
| for bb, sc in zip(b, s): |
| bw = bb[2] - bb[0] |
| bh = bb[3] - bb[1] |
| m = max(8, int(min(bw, bh) * 0.2)) |
|
|
| if not left_edge and bb[0] < m: |
| continue |
| if not top_edge and bb[1] < m: |
| continue |
| if not right_edge and bb[2] > (vw - m): |
| continue |
| if not bottom_edge and bb[3] > (vh - m): |
| continue |
|
|
| x1 = max(0, min(W - 1, int(bb[0] + ox))) |
| y1 = max(0, min(H - 1, int(bb[1] + oy))) |
| x2 = max(0, min(W - 1, int(bb[2] + ox))) |
| y2 = max(0, min(H - 1, int(bb[3] + oy))) |
|
|
| if x2 > x1 and y2 > y1: |
| all_boxes.append([x1, y1, x2, y2]) |
| all_scores.append(float(sc)) |
|
|
| if not all_boxes: |
| return [] |
|
|
| boxes = np.asarray(all_boxes, dtype=np.float32) |
| scores = np.asarray(all_scores, dtype=np.float32) |
|
|
| return self._merge_boxes_global(image.shape, boxes, scores) |
|
|
| def predict_batch( |
| self, |
| batch_images: list[ndarray], |
| offset: int, |
| n_keypoints: int, |
| ) -> list[TVFrameResult]: |
| results: list[TVFrameResult] = [] |
|
|
| for frame_number_in_batch, image in enumerate(batch_images): |
| try: |
| boxes = self.predict_image(image) |
| except Exception as e: |
| print(f"⚠️ Inference failed for frame {offset + frame_number_in_batch}: {e}") |
| boxes = [] |
|
|
| results.append( |
| TVFrameResult( |
| frame_id=offset + frame_number_in_batch, |
| boxes=boxes, |
| keypoints=[(0, 0) for _ in range(max(0, int(n_keypoints)))], |
| ) |
| ) |
|
|
| return results |
