from pathlib import Path import math import cv2 import numpy as np import 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: """ONNX Runtime miner with per-class candidates, TTA fusion, and temporal rescue.""" class_names = ["cup", "bottle", "can"] model_class_names = ["cup", "bottle", "can"] _model_to_competition_cls = np.array([0, 1, 2], dtype=np.int32) input_size = 1280 iou_thres = 0.3 cross_iou_thresh = 0.65 min_side = 8.0 min_box_area = 100.0 max_aspect_ratio = 10.0 max_det = 300 _conf_thres_array = np.array([0.60, 0.45, 0.50], dtype=np.float32) _candidate_conf_thres_array = np.array([0.20, 0.30, 0.30], dtype=np.float32) _tta_conf_thres_array = np.array([0.52, 0.37, 0.42], dtype=np.float32) _temporal_conf_thres_array = np.array([0.54, 0.39, 0.44], dtype=np.float32) _tta_confirmed_views = 1 temporal_iou_thresh = 0.25 track_iou_thresh = 0.35 track_keep_frames = 2 track_min_conf = np.array([0.50, 0.35, 0.40], dtype=np.float32) sparse_candidate_count = 8 crowded_candidate_count = 28 crowded_area_ratio = 0.030 sparse_relax = 0.04 crowded_raise = 0.04 def __init__(self, path_hf_repo: Path) -> None: model_path = path_hf_repo / "weights.onnx" print("ORT version:", ort.__version__) try: ort.preload_dlls() print("preload_dlls success") except Exception as e: print(f"preload_dlls failed: {e}") print("ORT available providers BEFORE session:", ort.get_available_providers()) 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"], ) print("Created ORT session with preferred CUDA provider list") except Exception as e: print(f"CUDA session creation failed, falling back to CPU: {e}") self.session = ort.InferenceSession( str(model_path), sess_options=sess_options, providers=["CPUExecutionProvider"], ) print("ORT session providers:", self.session.get_providers()) for inp in self.session.get_inputs(): print("INPUT:", inp.name, inp.shape, inp.type) for out in self.session.get_outputs(): print("OUTPUT:", out.name, out.shape, out.type) self.input_name = self.session.get_inputs()[0].name self.output_names = [output.name for output in self.session.get_outputs()] self.input_shape = self.session.get_inputs()[0].shape self.input_dtype = self._input_dtype(self.session.get_inputs()[0].type) self.input_height = self._safe_dim(self.input_shape[2], default=self.input_size) self.input_width = self._safe_dim(self.input_shape[3], default=self.input_size) self._tracks: list[dict] = [] self._last_track_frame_id: int | None = None print(f"ONNX model loaded from: {model_path}") print(f"ONNX providers: {self.session.get_providers()}") print(f"ONNX input: name={self.input_name}, shape={self.input_shape}") print(f"ONNX input dtype: {self.input_dtype}") 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 _input_dtype(input_type: str) -> np.dtype: if input_type == "tensor(float16)": return np.dtype(np.float16) return np.dtype(np.float32) @classmethod def _to_competition_cls(cls, cls_ids: np.ndarray) -> np.ndarray: if len(cls_ids) == 0: return cls_ids.astype(np.int32) valid = (cls_ids >= 0) & (cls_ids < len(cls._model_to_competition_cls)) remapped = np.full_like(cls_ids, fill_value=-1, dtype=np.int32) remapped[valid] = cls._model_to_competition_cls[cls_ids[valid]] return remapped def _letterbox(self, image: ndarray, new_shape: tuple[int, int], color=(114, 114, 114) ) -> tuple[ndarray, float, tuple[float, float]]: h, w = image.shape[:2] new_w, new_h = new_shape ratio = min(new_w / w, new_h / h) resized_w = int(round(w * ratio)) resized_h = int(round(h * ratio)) if (resized_w, resized_h) != (w, h): interp = cv2.INTER_CUBIC if ratio > 1.0 else cv2.INTER_LINEAR image = cv2.resize(image, (resized_w, resized_h), interpolation=interp) dw = (new_w - resized_w) / 2.0 dh = (new_h - resized_h) / 2.0 left = int(round(dw - 0.1)) right = int(round(dw + 0.1)) top = int(round(dh - 0.1)) bottom = int(round(dh + 0.1)) padded = cv2.copyMakeBorder(image, top, bottom, left, right, borderType=cv2.BORDER_CONSTANT, value=color) return padded, ratio, (dw, dh) def _preprocess(self, image: ndarray ) -> tuple[np.ndarray, float, tuple[float, float], tuple[int, int]]: orig_h, orig_w = image.shape[:2] img, ratio, pad = self._letterbox(image, (self.input_width, self.input_height)) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = img.astype(self.input_dtype) / self.input_dtype.type(255.0) img = np.transpose(img, (2, 0, 1))[None, ...] img = np.ascontiguousarray(img, dtype=self.input_dtype) return img, ratio, pad, (orig_w, orig_h) @staticmethod def _clip_boxes(boxes: np.ndarray, image_size: tuple[int, int]) -> np.ndarray: w, h = image_size 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 @staticmethod def _hard_nms(boxes: np.ndarray, scores: np.ndarray, iou_thresh: float) -> np.ndarray: n = len(boxes) if n == 0: return np.array([], dtype=np.intp) order = np.argsort(-scores) keep: list[int] = [] while len(order) > 0: i = int(order[0]) keep.append(i) if len(order) == 1: break rest = order[1:] xx1 = np.maximum(boxes[i, 0], boxes[rest, 0]) yy1 = np.maximum(boxes[i, 1], boxes[rest, 1]) xx2 = np.minimum(boxes[i, 2], boxes[rest, 2]) yy2 = np.minimum(boxes[i, 3], boxes[rest, 3]) inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1) a_i = (max(0.0, boxes[i, 2] - boxes[i, 0]) * max(0.0, boxes[i, 3] - boxes[i, 1])) a_r = (np.maximum(0.0, boxes[rest, 2] - boxes[rest, 0]) * np.maximum(0.0, boxes[rest, 3] - boxes[rest, 1])) iou = inter / (a_i + a_r - inter + 1e-7) order = rest[iou <= iou_thresh] return np.array(keep, dtype=np.intp) def _per_class_hard_nms(self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray, iou_thresh: float ) -> np.ndarray: if len(boxes) == 0: return np.array([], dtype=np.intp) all_keep: list[int] = [] for c in np.unique(cls_ids): mask = cls_ids == c indices = np.where(mask)[0] keep = self._hard_nms(boxes[mask], scores[mask], iou_thresh) all_keep.extend(indices[keep].tolist()) all_keep.sort() return np.array(all_keep, dtype=np.intp) def _cross_class_dedup_op(self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray, iou_thresh: float ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: keep_idx = self._cross_class_dedup_keep_indices( boxes, scores, cls_ids, iou_thresh ) return boxes[keep_idx], scores[keep_idx], cls_ids[keep_idx] def _cross_class_dedup_keep_indices(self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray, iou_thresh: float) -> np.ndarray: n = len(boxes) if n <= 1: return np.arange(n, dtype=np.intp) boxes = np.asarray(boxes, dtype=np.float32) scores = np.asarray(scores, dtype=np.float32) cls_ids = np.asarray(cls_ids, dtype=np.int32) areas = (np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) * np.maximum(0.0, boxes[:, 3] - boxes[:, 1])) margins = scores - self._conf_thres_array[cls_ids] order = np.lexsort((-areas, -margins)) suppressed = np.zeros(n, dtype=bool) keep: list[int] = [] for i in order: if suppressed[i]: continue keep.append(int(i)) bi = boxes[i] xx1 = np.maximum(bi[0], boxes[:, 0]) yy1 = np.maximum(bi[1], boxes[:, 1]) xx2 = np.minimum(bi[2], boxes[:, 2]) yy2 = np.minimum(bi[3], boxes[:, 3]) inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1) a_i = max(1e-7, float((bi[2] - bi[0]) * (bi[3] - bi[1]))) iou = inter / (a_i + areas - inter + 1e-7) dup = iou > iou_thresh dup[i] = False suppressed |= dup return np.array(keep, dtype=np.intp) def _filter_sane_boxes(self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray, orig_size: tuple[int, int] ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: if len(boxes) == 0: return boxes, scores, cls_ids orig_w, orig_h = orig_size image_area = float(orig_w * orig_h) bw = np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) bh = np.maximum(0.0, boxes[:, 3] - boxes[:, 1]) area = bw * bh ar = np.where( (bw > 0) & (bh > 0), np.maximum(bw / np.maximum(bh, 1e-6), bh / np.maximum(bw, 1e-6)), np.inf, ) keep = ( (bw >= self.min_side) & (bh >= self.min_side) & (area >= self.min_box_area) & (area <= 0.95 * image_area) & (ar <= self.max_aspect_ratio) ) return boxes[keep], scores[keep], cls_ids[keep] def _max_score_per_cluster(self, post_boxes: np.ndarray, post_cls: np.ndarray, full_boxes: np.ndarray, full_scores: np.ndarray, full_cls: np.ndarray, iou_thresh: float) -> np.ndarray: n = len(post_boxes) if n == 0: return np.empty(0, dtype=np.float32) full_areas = (np.maximum(0.0, full_boxes[:, 2] - full_boxes[:, 0]) * np.maximum(0.0, full_boxes[:, 3] - full_boxes[:, 1])) out = np.empty(n, dtype=np.float32) for i in range(n): bi = post_boxes[i] xx1 = np.maximum(bi[0], full_boxes[:, 0]) yy1 = np.maximum(bi[1], full_boxes[:, 1]) xx2 = np.minimum(bi[2], full_boxes[:, 2]) yy2 = np.minimum(bi[3], full_boxes[:, 3]) inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1) a_i = max(0.0, float((bi[2] - bi[0]) * (bi[3] - bi[1]))) iou = inter / (a_i + full_areas - inter + 1e-7) cluster = (iou >= iou_thresh) & (full_cls == post_cls[i]) out[i] = float(np.max(full_scores[cluster])) if np.any(cluster) else 0.0 return out def _view_support_per_cluster(self, post_boxes: np.ndarray, post_cls: np.ndarray, full_boxes: np.ndarray, full_cls: np.ndarray, full_view_ids: np.ndarray, iou_thresh: float) -> np.ndarray: n = len(post_boxes) if n == 0: return np.empty(0, dtype=np.int32) full_areas = (np.maximum(0.0, full_boxes[:, 2] - full_boxes[:, 0]) * np.maximum(0.0, full_boxes[:, 3] - full_boxes[:, 1])) out = np.ones(n, dtype=np.int32) for i in range(n): bi = post_boxes[i] xx1 = np.maximum(bi[0], full_boxes[:, 0]) yy1 = np.maximum(bi[1], full_boxes[:, 1]) xx2 = np.minimum(bi[2], full_boxes[:, 2]) yy2 = np.minimum(bi[3], full_boxes[:, 3]) inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1) a_i = max(0.0, float((bi[2] - bi[0]) * (bi[3] - bi[1]))) iou = inter / (a_i + full_areas - inter + 1e-7) cluster = (iou >= iou_thresh) & (full_cls == post_cls[i]) if np.any(cluster): out[i] = int(len(np.unique(full_view_ids[cluster]))) return out def _conf_filter_mask(self, scores: np.ndarray, cls_ids: np.ndarray) -> np.ndarray: """Keep low-score candidates; final acceptance happens after evidence fusion.""" if len(scores) == 0: return np.zeros(0, dtype=bool) return scores >= self._candidate_conf_thres_array[cls_ids] def _scene_adjustment(self, boxes: list[BoundingBox], image_shape: tuple[int, int, int] | None) -> float: if not boxes or image_shape is None: return 0.0 h, w = image_shape[:2] image_area = max(1.0, float(w * h)) total_box_area = sum( max(0, box.x2 - box.x1) * max(0, box.y2 - box.y1) for box in boxes ) area_ratio = float(total_box_area) / image_area if len(boxes) >= self.crowded_candidate_count or area_ratio >= self.crowded_area_ratio: return self.crowded_raise if len(boxes) <= self.sparse_candidate_count and area_ratio < self.crowded_area_ratio * 0.5: return -self.sparse_relax return 0.0 def _adaptive_thresholds(self, cls_ids: np.ndarray, boxes: list[BoundingBox], image_shape: tuple[int, int, int] | None ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: adjustment = self._scene_adjustment(boxes, image_shape) auto = np.clip( self._conf_thres_array[cls_ids] + adjustment, self._candidate_conf_thres_array[cls_ids] + 0.05, 0.95, ) tta = np.clip( self._tta_conf_thres_array[cls_ids] + adjustment, self._candidate_conf_thres_array[cls_ids], auto, ) temporal = np.clip( self._temporal_conf_thres_array[cls_ids] + adjustment, self._candidate_conf_thres_array[cls_ids], auto, ) return auto, tta, temporal def _candidate_accept_mask(self, boxes: list[BoundingBox], view_support: np.ndarray, image_shape: tuple[int, int, int] | None ) -> np.ndarray: _, scores, cls_ids = self._boxes_to_arrays(boxes) if len(scores) == 0: return np.zeros(0, dtype=bool) auto_thres, tta_thres, _ = self._adaptive_thresholds( cls_ids, boxes, image_shape ) auto = scores >= auto_thres tta_confirmed = ( (scores >= tta_thres) & (view_support >= self._tta_confirmed_views) ) return auto | tta_confirmed @staticmethod def _boxes_to_arrays(boxes: list[BoundingBox] ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: if not boxes: return ( np.empty((0, 4), dtype=np.float32), np.empty(0, dtype=np.float32), np.empty(0, dtype=np.int32), ) coords = np.array( [[b.x1, b.y1, b.x2, b.y2] for b in boxes], dtype=np.float32 ) scores = np.array([b.conf for b in boxes], dtype=np.float32) cls_ids = np.array([b.cls_id for b in boxes], dtype=np.int32) return coords, scores, cls_ids @staticmethod def _image_shape(image: np.ndarray | None) -> tuple[int, int, int] | None: if isinstance(image, np.ndarray) and image.ndim == 3: return image.shape return None @staticmethod def _single_box_iou(box: BoundingBox, boxes: np.ndarray) -> np.ndarray: if len(boxes) == 0: return np.empty(0, dtype=np.float32) xx1 = np.maximum(float(box.x1), boxes[:, 0]) yy1 = np.maximum(float(box.y1), boxes[:, 1]) xx2 = np.minimum(float(box.x2), boxes[:, 2]) yy2 = np.minimum(float(box.y2), boxes[:, 3]) inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1) a_i = max(0.0, float((box.x2 - box.x1) * (box.y2 - box.y1))) areas = (np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) * np.maximum(0.0, boxes[:, 3] - boxes[:, 1])) return inter / (a_i + areas - inter + 1e-7) def _per_view_pipeline(self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray, orig_size: tuple[int, int] ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: boxes, scores, cls_ids = self._filter_sane_boxes( boxes, scores, cls_ids, orig_size ) if len(boxes) == 0: return boxes, scores, cls_ids if len(boxes) > 1: keep = self._per_class_hard_nms(boxes, scores, cls_ids, self.iou_thres) boxes, scores, cls_ids = boxes[keep], scores[keep], cls_ids[keep] if len(scores) > self.max_det: top = np.argsort(-scores)[: self.max_det] boxes, scores, cls_ids = boxes[top], scores[top], cls_ids[top] if len(boxes) > 1: boxes, scores, cls_ids = self._cross_class_dedup_op( boxes, scores, cls_ids, self.cross_iou_thresh ) return boxes, scores, cls_ids def _decode_final_dets(self, preds: np.ndarray, ratio: float, pad: tuple[float, float], orig_size: tuple[int, int]) -> list[BoundingBox]: 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 ONNX final-det output shape: {preds.shape}") boxes = preds[:, :4].astype(np.float32) scores = preds[:, 4].astype(np.float32) cls_ids = self._to_competition_cls(preds[:, 5].astype(np.int32)) valid_cls = cls_ids >= 0 boxes = boxes[valid_cls] scores = scores[valid_cls] cls_ids = cls_ids[valid_cls] keep = self._conf_filter_mask(scores, cls_ids) boxes = boxes[keep] scores = scores[keep] cls_ids = cls_ids[keep] if len(boxes) == 0: return [] pad_w, pad_h = pad boxes[:, [0, 2]] -= pad_w boxes[:, [1, 3]] -= pad_h boxes /= ratio boxes = self._clip_boxes(boxes, orig_size) boxes, scores, cls_ids = self._per_view_pipeline( boxes, scores, cls_ids, orig_size ) return self._build_results(boxes, scores, cls_ids) def _decode_raw_yolo(self, preds: np.ndarray, ratio: float, pad: tuple[float, float], orig_size: tuple[int, int]) -> list[BoundingBox]: if preds.ndim != 3 or preds.shape[0] != 1: raise ValueError(f"Unexpected raw ONNX output shape: {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 raw output shape: {preds.shape}") boxes_xywh = preds[:, :4].astype(np.float32) cls_part = preds[:, 4:].astype(np.float32) if cls_part.shape[1] == 1: scores = cls_part[:, 0] cls_ids = np.zeros(len(scores), dtype=np.int32) else: cls_ids = np.argmax(cls_part, axis=1).astype(np.int32) scores = cls_part[np.arange(len(cls_part)), cls_ids] cls_ids = self._to_competition_cls(cls_ids) valid_cls = cls_ids >= 0 boxes_xywh = boxes_xywh[valid_cls] scores = scores[valid_cls] cls_ids = cls_ids[valid_cls] keep = self._conf_filter_mask(scores, cls_ids) boxes_xywh = boxes_xywh[keep] scores = scores[keep] cls_ids = cls_ids[keep] if len(boxes_xywh) == 0: return [] boxes = self._xywh_to_xyxy(boxes_xywh) pad_w, pad_h = pad boxes[:, [0, 2]] -= pad_w boxes[:, [1, 3]] -= pad_h boxes /= ratio boxes = self._clip_boxes(boxes, orig_size) boxes, scores, cls_ids = self._per_view_pipeline( boxes, scores, cls_ids, orig_size ) return self._build_results(boxes, scores, cls_ids) @staticmethod def _build_results(boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray) -> list[BoundingBox]: results: list[BoundingBox] = [] for box, conf, cls_id in zip(boxes, scores, cls_ids): x1, y1, x2, y2 = box.tolist() if x2 <= x1 or y2 <= y1: continue results.append( BoundingBox( x1=int(math.floor(x1)), y1=int(math.floor(y1)), x2=int(math.ceil(x2)), y2=int(math.ceil(y2)), cls_id=int(cls_id), conf=float(conf), ) ) return results def _postprocess(self, output: np.ndarray, ratio: float, pad: tuple[float, float], orig_size: tuple[int, int]) -> list[BoundingBox]: if output.ndim == 2 and output.shape[1] >= 6: return self._decode_final_dets(output, ratio, pad, orig_size) if output.ndim == 3 and output.shape[0] == 1 and output.shape[2] == 6: return self._decode_final_dets(output, ratio, pad, orig_size) return self._decode_raw_yolo(output, ratio, pad, orig_size) def _predict_single(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}") if image.dtype != np.uint8: image = image.astype(np.uint8) input_tensor, ratio, pad, orig_size = self._preprocess(image) expected = (1, 3, self.input_height, self.input_width) if input_tensor.shape != expected: raise ValueError( f"Bad input tensor shape={input_tensor.shape}, expected={expected}" ) outputs = self.session.run(self.output_names, {self.input_name: input_tensor}) return self._postprocess(outputs[0], ratio, pad, orig_size) def _predict_tta_candidates(self, image: np.ndarray ) -> tuple[list[BoundingBox], np.ndarray]: boxes_orig = self._predict_single(image) flipped = cv2.flip(image, 1) boxes_flip = self._predict_single(flipped) w = image.shape[1] boxes_flip = [ BoundingBox( x1=w - b.x2, y1=b.y1, x2=w - b.x1, y2=b.y2, cls_id=b.cls_id, conf=b.conf, ) for b in boxes_flip ] all_boxes = boxes_orig + boxes_flip if not all_boxes: return [], np.empty(0, dtype=np.int32) coords, scores, cls_ids = self._boxes_to_arrays(all_boxes) view_ids = np.array( [0] * len(boxes_orig) + [1] * len(boxes_flip), dtype=np.int32 ) hard_keep = self._per_class_hard_nms(coords, scores, cls_ids, self.iou_thres) if len(hard_keep) == 0: return [], np.empty(0, dtype=np.int32) if len(hard_keep) > self.max_det: top = np.argsort(-scores[hard_keep])[: self.max_det] hard_keep = hard_keep[top] boosted = self._max_score_per_cluster( coords[hard_keep], cls_ids[hard_keep], coords, scores, cls_ids, self.iou_thres, ) kept_coords = coords[hard_keep] kept_cls = cls_ids[hard_keep] view_support = self._view_support_per_cluster( kept_coords, kept_cls, coords, cls_ids, view_ids, self.iou_thres, ) if len(kept_coords) > 1: dedup_keep = self._cross_class_dedup_keep_indices( kept_coords, boosted, kept_cls, self.cross_iou_thresh ) kept_coords = kept_coords[dedup_keep] boosted = boosted[dedup_keep] kept_cls = kept_cls[dedup_keep] view_support = view_support[dedup_keep] boxes = [ BoundingBox( x1=int(math.floor(kept_coords[j, 0])), y1=int(math.floor(kept_coords[j, 1])), x2=int(math.ceil(kept_coords[j, 2])), y2=int(math.ceil(kept_coords[j, 3])), cls_id=int(kept_cls[j]), conf=float(boosted[j]), ) for j in range(len(kept_coords)) ] return boxes, view_support def _filter_by_evidence(self, boxes: list[BoundingBox], view_support: np.ndarray, image_shape: tuple[int, int, int] | None ) -> list[BoundingBox]: keep = self._candidate_accept_mask(boxes, view_support, image_shape) return [box for box, ok in zip(boxes, keep) if bool(ok)] def _predict_tta(self, image: np.ndarray) -> list[BoundingBox]: boxes, view_support = self._predict_tta_candidates(image) return self._filter_by_evidence(boxes, view_support, image.shape) def _has_temporal_support(self, frame_idx: int, box: BoundingBox, candidate_boxes: list[list[BoundingBox]], initial_keep: list[np.ndarray]) -> bool: neighbor_indices = [ idx for idx in (frame_idx - 1, frame_idx + 1) if 0 <= idx < len(candidate_boxes) ] for idx in neighbor_indices: coords, _, cls_ids = self._boxes_to_arrays(candidate_boxes[idx]) same_cls = cls_ids == box.cls_id if not np.any(same_cls): continue accepted = same_cls & initial_keep[idx] if np.any(accepted): if np.max(self._single_box_iou(box, coords[accepted])) >= self.temporal_iou_thresh: return True two_sided_candidate_support = [] for idx in (frame_idx - 1, frame_idx + 1): if not 0 <= idx < len(candidate_boxes): two_sided_candidate_support.append(False) continue coords, scores, cls_ids = self._boxes_to_arrays(candidate_boxes[idx]) same_cls = cls_ids == box.cls_id if not np.any(same_cls): two_sided_candidate_support.append(False) continue score_ok = scores >= self._temporal_conf_thres_array[cls_ids] neighbor_ok = same_cls & score_ok supported = ( np.any(neighbor_ok) and np.max(self._single_box_iou(box, coords[neighbor_ok])) >= self.temporal_iou_thresh ) two_sided_candidate_support.append(bool(supported)) return all(two_sided_candidate_support) def _reset_tracks_if_needed(self, frame_id: int) -> None: if self._last_track_frame_id is None: self._last_track_frame_id = frame_id - 1 return if frame_id <= self._last_track_frame_id: self._tracks = [] self._last_track_frame_id = frame_id def _track_supported(self, box: BoundingBox, frame_id: int) -> bool: best_iou = 0.0 for track in self._tracks: if int(track["cls_id"]) != box.cls_id: continue age = frame_id - int(track["frame_id"]) if age < 1 or age > self.track_keep_frames: continue iou = self._single_box_iou(box, track["coords"])[0] best_iou = max(best_iou, float(iou)) return best_iou >= self.track_iou_thresh def _update_tracks(self, boxes: list[BoundingBox], frame_id: int) -> None: fresh_tracks = [] for track in self._tracks: if frame_id - int(track["frame_id"]) <= self.track_keep_frames: fresh_tracks.append(track) for box in boxes: coords = np.array( [[box.x1, box.y1, box.x2, box.y2]], dtype=np.float32 ) updated = False for track in fresh_tracks: if int(track["cls_id"]) != box.cls_id: continue iou = self._single_box_iou(box, track["coords"])[0] if iou >= self.track_iou_thresh: track["coords"] = coords track["frame_id"] = frame_id track["conf"] = box.conf updated = True break if not updated: fresh_tracks.append( { "coords": coords, "cls_id": box.cls_id, "conf": box.conf, "frame_id": frame_id, } ) self._tracks = fresh_tracks def predict_batch(self, batch_images: list[ndarray], offset: int, n_keypoints: int) -> list[TVFrameResult]: candidate_boxes: list[list[BoundingBox]] = [] view_supports: list[np.ndarray] = [] image_shapes = [self._image_shape(image) for image in batch_images] results: list[TVFrameResult] = [] for frame_number_in_batch, image in enumerate(batch_images): try: boxes, view_support = self._predict_tta_candidates(image) except Exception as e: print(f"Inference failed for frame {offset + frame_number_in_batch}: {e}") boxes = [] view_support = np.empty(0, dtype=np.int32) candidate_boxes.append(boxes) view_supports.append(view_support) initial_keep: list[np.ndarray] = [] for boxes, view_support, image_shape in zip( candidate_boxes, view_supports, image_shapes ): initial_keep.append( self._candidate_accept_mask(boxes, view_support, image_shape) ) for frame_number_in_batch, boxes in enumerate(candidate_boxes): frame_id = offset + frame_number_in_batch self._reset_tracks_if_needed(frame_id) keep = initial_keep[frame_number_in_batch].copy() _, scores, cls_ids = self._boxes_to_arrays(boxes) _, _, temporal_thres = self._adaptive_thresholds( cls_ids, boxes, image_shapes[frame_number_in_batch] ) temporal_ready = scores >= temporal_thres track_ready = scores >= self.track_min_conf[cls_ids] for i, box in enumerate(boxes): if keep[i]: continue has_neighbor_support = ( bool(temporal_ready[i]) and self._has_temporal_support( frame_number_in_batch, box, candidate_boxes, initial_keep ) ) has_track_support = ( bool(track_ready[i]) and self._track_supported(box, frame_id) ) if has_neighbor_support or has_track_support: keep[i] = True boxes = [box for box, ok in zip(boxes, keep) if bool(ok)] self._update_tracks(boxes, frame_id) results.append( TVFrameResult( frame_id=frame_id, boxes=boxes, keypoints=[(0, 0) for _ in range(max(0, int(n_keypoints)))], ) ) return results # if __name__ == "__main__": # # predict batch with local images for testing # import json # from time import time # test_images = [ # cv2.imread(str(p)) for p in sorted(Path("test_images").glob("*.jpg")) # ] # miner = Miner(Path("hf_repo")) # start_time = time() # results = miner.predict_batch(test_images, offset=0, n_keypoints=0) # end_time = time() # print(f"Predicted {len(test_images)} images in {end_time - start_time:.2f} seconds") # print(json.dumps([r.dict() for r in results], indent=2))