Martin99113 commited on
Commit
d9cd193
·
verified ·
1 Parent(s): 01f852b

scorevision: push artifact

Browse files
Files changed (1) hide show
  1. miner.py +1059 -0
miner.py ADDED
@@ -0,0 +1,1059 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from pathlib import Path
2
+ import math
3
+
4
+ import cv2
5
+ import numpy as np
6
+ import onnxruntime as ort
7
+ from numpy import ndarray
8
+ from pydantic import BaseModel
9
+
10
+ class BoundingBox(BaseModel):
11
+ x1: int
12
+ y1: int
13
+ x2: int
14
+ y2: int
15
+ cls_id: int
16
+ conf: float
17
+
18
+
19
+ class TVFrameResult(BaseModel):
20
+ frame_id: int
21
+ boxes: list[BoundingBox]
22
+ keypoints: list[tuple[int, int]]
23
+
24
+
25
+ class Miner:
26
+ """ONNX Runtime miner for road-sign detection (single class).
27
+
28
+ Post-processing: per-class confidence + rescue bonus, hard NMS,
29
+ cross-class dedup, sanity-box filter, optional TTA / tile merge.
30
+ """
31
+
32
+ class_names = ["road_sign"]
33
+ _model_class_order = ["road_sign"]
34
+
35
+ iou_thres = 0.8
36
+ cross_iou_thresh = 0.8
37
+ max_det = 150
38
+
39
+ use_secondary_merge = True
40
+ secondary_conf = 0.80
41
+ merge_iou = 0.2
42
+ dual_head_dedup_iou = 0.35
43
+ remove_contained_boxes = True
44
+ _conf_thres_array = np.array(
45
+ [0.37], dtype=np.float32
46
+ )
47
+ _bonus_array = np.array(
48
+ [0.2], dtype=np.float32
49
+ )
50
+
51
+ min_box_area = 4 * 4
52
+ min_side = 3
53
+ max_aspect_ratio = 12.0
54
+
55
+ tile_trigger_ratio = 1.4
56
+ tile_overlap_ratio = 0.20
57
+ # use_fast_postprocess = True
58
+
59
+ @staticmethod
60
+ def _ort_provider_chain() -> list[str]:
61
+ """Prefer GPU/DML providers when installed; always end with CPU fallback."""
62
+ available = set(ort.get_available_providers())
63
+ preferred = (
64
+ "CPUExecutionProvider",
65
+ "CUDAExecutionProvider",
66
+ "DmlExecutionProvider",
67
+ )
68
+ return [p for p in preferred if p in available] or ["CPUExecutionProvider"]
69
+
70
+ @staticmethod
71
+ def _ort_session_options() -> ort.SessionOptions:
72
+ sess_options = ort.SessionOptions()
73
+ sess_options.enable_cpu_mem_arena = True
74
+ sess_options.enable_mem_pattern = True
75
+ sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
76
+ sess_options.intra_op_num_threads = 2
77
+ sess_options.inter_op_num_threads = 1
78
+ sess_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
79
+ return sess_options
80
+
81
+ def __init__(self, path_hf_repo: Path) -> None:
82
+ model_path = path_hf_repo / "weights.onnx"
83
+ print("ORT version:", ort.__version__)
84
+
85
+ try:
86
+ ort.preload_dlls()
87
+ print("✅ onnxruntime.preload_dlls() success")
88
+ except Exception as e:
89
+ print(f"⚠️ preload_dlls failed: {e}")
90
+
91
+ print("ORT available providers BEFORE session:", ort.get_available_providers())
92
+
93
+ sess_options = self._ort_session_options()
94
+ providers = self._ort_provider_chain()
95
+ print(
96
+ f"ORT session config: providers={providers} "
97
+ f"intra_op_threads={sess_options.intra_op_num_threads}"
98
+ )
99
+ self.session = ort.InferenceSession(
100
+ str(model_path),
101
+ sess_options=sess_options,
102
+ providers=providers,
103
+ )
104
+
105
+ print("ORT session providers:", self.session.get_providers())
106
+
107
+ # Build cls_remap: for each model-emit index i,
108
+ # cls_remap[i] = self.class_names.index(model_class_order[i])
109
+ # i.e. convert a model-side class id into the output class id that
110
+ # downstream code (BoundingBox.cls_id, the per-class threshold/bonus
111
+ # arrays) expects. The model-side order comes from the ONNX metadata
112
+ # when available, else falls back to the static _model_class_order.
113
+ model_class_order = self._read_model_class_order()
114
+ if model_class_order is None:
115
+ model_class_order = list(self._model_class_order)
116
+ print(f"cls order: no usable ONNX metadata, FALLBACK {model_class_order}")
117
+ else:
118
+ print(f"cls order: from ONNX metadata {model_class_order}")
119
+ self.cls_remap = np.array(
120
+ [self.class_names.index(n) for n in model_class_order],
121
+ dtype=np.int32,
122
+ )
123
+
124
+ for i, inp in enumerate(self.session.get_inputs()):
125
+ print(f"INPUT[{i}]: shape={inp.shape} type={inp.type}")
126
+ for i, out in enumerate(self.session.get_outputs()):
127
+ print(f"OUTPUT[{i}]: shape={out.shape} type={out.type}")
128
+
129
+ self.input_name = self.session.get_inputs()[0].name
130
+ self.output_names = [output.name for output in self.session.get_outputs()]
131
+ self.input_shape = self.session.get_inputs()[0].shape
132
+
133
+ # weights.onnx is exported at 1280x1280 (Ultralytics imgsz metadata),
134
+ # static (dynamic=False). The default is only the fallback for when the
135
+ # ONNX input dims aren't fixed; the real value is read from the session.
136
+ self.input_height = self._safe_dim(self.input_shape[2], default=1280)
137
+ self.input_width = self._safe_dim(self.input_shape[3], default=1280)
138
+
139
+ self.use_tta = False
140
+ self.use_tile_tta = False
141
+ # Soft-NMS (ported from carwash001): Gaussian score decay of overlapping
142
+ # boxes instead of hard removal. OFF by default to preserve the current
143
+ # deployed behaviour; flip on (and tune sigma) via tune_miner.py to see if
144
+ # it scores better — useful where signs cluster (gantries, sign assemblies).
145
+ self.use_soft_nms = False
146
+ self.soft_nms_sigma = 0.5
147
+ self.soft_nms_score_thresh = 0.01
148
+
149
+ print(f"✅ ONNX model loaded from: {model_path}")
150
+ print(f"✅ ONNX providers: {self.session.get_providers()}")
151
+ print(f"✅ ONNX input shape={self.input_shape}")
152
+ print(f"✅ ONNX input size: {self.input_width}x{self.input_height}, "
153
+ f"use_tta={self.use_tta}, use_tile_tta={self.use_tile_tta}")
154
+ print("per-class conf: " + ", ".join(
155
+ f"{n}={t:.3f}" for n, t in zip(
156
+ self.class_names, self._conf_thres_array.tolist()
157
+ )
158
+ ))
159
+
160
+ self._warmup()
161
+
162
+ def _warmup(self, iters: int = 3) -> None:
163
+ try:
164
+ dummy = np.zeros((720, 1280, 3), dtype=np.uint8)
165
+ for _ in range(max(1, iters)):
166
+ self.predict_batch(batch_images=[dummy], offset=0, n_keypoints=0)
167
+ print(f"✅ warmup: {iters} dummy predict_batch call(s) done")
168
+ except Exception as e:
169
+ print(f"⚠️ warmup skipped: {e}")
170
+
171
+ def _read_model_class_order(self) -> "list[str] | None":
172
+ """Read the model's class order from Ultralytics ONNX metadata.
173
+
174
+ Returns the class names ordered by model-emit index, or None when the
175
+ metadata is missing/unparsable or doesn't match `class_names` as a set
176
+ (in which case the static _model_class_order fallback is used)."""
177
+ try:
178
+ import ast
179
+
180
+ meta = self.session.get_modelmeta().custom_metadata_map
181
+ names = ast.literal_eval(meta["names"]) # e.g. {0: 'road_sign'}
182
+ if isinstance(names, dict):
183
+ order = [str(names[i]) for i in sorted(names)]
184
+ else:
185
+ order = [str(n) for n in names]
186
+ except Exception as e:
187
+ print(f"cls order: could not read ONNX names metadata ({e})")
188
+ return None
189
+ if sorted(order) != sorted(self.class_names):
190
+ print(
191
+ f"cls order: ONNX names {order} do not match expected classes "
192
+ f"{self.class_names}; ignoring metadata"
193
+ )
194
+ return None
195
+ return order
196
+
197
+ def __repr__(self) -> str:
198
+ return (
199
+ f"ONNXRuntime(session={type(self.session).__name__}, "
200
+ f"providers={self.session.get_providers()})"
201
+ )
202
+
203
+ @staticmethod
204
+ def _safe_dim(value, default: int) -> int:
205
+ return value if isinstance(value, int) and value > 0 else default
206
+
207
+ def _letterbox(
208
+ self,
209
+ image: ndarray,
210
+ new_shape: tuple[int, int],
211
+ color=(114, 114, 114),
212
+ ) -> tuple[ndarray, float, tuple[float, float]]:
213
+ h, w = image.shape[:2]
214
+ new_w, new_h = new_shape
215
+
216
+ ratio = min(new_w / w, new_h / h)
217
+ resized_w = int(round(w * ratio))
218
+ resized_h = int(round(h * ratio))
219
+
220
+ if (resized_w, resized_h) != (w, h):
221
+ interp = cv2.INTER_CUBIC if ratio > 1.0 else cv2.INTER_LINEAR
222
+ image = cv2.resize(image, (resized_w, resized_h), interpolation=interp)
223
+
224
+ dw = (new_w - resized_w) / 2.0
225
+ dh = (new_h - resized_h) / 2.0
226
+
227
+ left = int(round(dw - 0.1))
228
+ right = int(round(dw + 0.1))
229
+ top = int(round(dh - 0.1))
230
+ bottom = int(round(dh + 0.1))
231
+
232
+ padded = cv2.copyMakeBorder(
233
+ image, top, bottom, left, right,
234
+ borderType=cv2.BORDER_CONSTANT, value=color,
235
+ )
236
+ return padded, ratio, (dw, dh)
237
+
238
+ def _preprocess(
239
+ self, image: ndarray
240
+ ) -> tuple[np.ndarray, float, tuple[float, float], tuple[int, int]]:
241
+ orig_h, orig_w = image.shape[:2]
242
+ img, ratio, pad = self._letterbox(
243
+ image, (self.input_width, self.input_height)
244
+ )
245
+ # Fused scale(1/255) + BGR->RGB swap + HWC->NCHW + contiguous float32 in
246
+ # one optimized OpenCV call (bit-identical to the cvtColor + astype/255 +
247
+ # transpose chain, but ~half the preprocess time).
248
+ blob = cv2.dnn.blobFromImage(img, scalefactor=1.0 / 255.0, swapRB=True)
249
+ return blob, ratio, pad, (orig_w, orig_h)
250
+
251
+ @staticmethod
252
+ def _clip_boxes(boxes: np.ndarray, image_size: tuple[int, int]) -> np.ndarray:
253
+ w, h = image_size
254
+ boxes[:, 0] = np.clip(boxes[:, 0], 0, w - 1)
255
+ boxes[:, 1] = np.clip(boxes[:, 1], 0, h - 1)
256
+ boxes[:, 2] = np.clip(boxes[:, 2], 0, w - 1)
257
+ boxes[:, 3] = np.clip(boxes[:, 3], 0, h - 1)
258
+ return boxes
259
+
260
+ @staticmethod
261
+ def _box_area(b: BoundingBox) -> int:
262
+ return max(0, b.x2 - b.x1) * max(0, b.y2 - b.y1)
263
+
264
+ @staticmethod
265
+ def _box_fully_contains(outer: BoundingBox, inner: BoundingBox) -> bool:
266
+ """True when outer strictly contains inner (outer area must be larger)."""
267
+ if (outer.x1 > inner.x1 or outer.y1 > inner.y1
268
+ or outer.x2 < inner.x2 or outer.y2 < inner.y2):
269
+ return False
270
+ outer_area = max(0, outer.x2 - outer.x1) * max(0, outer.y2 - outer.y1)
271
+ inner_area = max(0, inner.x2 - inner.x1) * max(0, inner.y2 - inner.y1)
272
+ return outer_area > inner_area
273
+
274
+ def _remove_contained_boxes(self, boxes: list[BoundingBox]) -> list[BoundingBox]:
275
+ """Drop smaller boxes fully contained inside a larger kept box."""
276
+ if not boxes or not self.remove_contained_boxes:
277
+ return boxes
278
+
279
+ sorted_boxes = sorted(boxes, key=self._box_area, reverse=True)
280
+ kept: list[BoundingBox] = []
281
+ for b in sorted_boxes:
282
+ if any(self._box_fully_contains(k, b) for k in kept):
283
+ continue
284
+ kept.append(b)
285
+ return kept
286
+
287
+ def _finalize_boxes(
288
+ self,
289
+ boxes: list[BoundingBox],
290
+ image_size: tuple[int, int],
291
+ ) -> list[BoundingBox]:
292
+ """Expand boxes, then remove smaller boxes contained in larger ones."""
293
+
294
+ return self._remove_contained_boxes(boxes)
295
+
296
+ @staticmethod
297
+ def _xywh_to_xyxy(boxes: np.ndarray) -> np.ndarray:
298
+ out = np.empty_like(boxes)
299
+ out[:, 0] = boxes[:, 0] - boxes[:, 2] / 2.0
300
+ out[:, 1] = boxes[:, 1] - boxes[:, 3] / 2.0
301
+ out[:, 2] = boxes[:, 0] + boxes[:, 2] / 2.0
302
+ out[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.0
303
+ return out
304
+
305
+ @staticmethod
306
+ def _hard_nms(
307
+ boxes: np.ndarray, scores: np.ndarray, iou_thresh: float
308
+ ) -> np.ndarray:
309
+ n = len(boxes)
310
+ if n == 0:
311
+ return np.array([], dtype=np.intp)
312
+ order = np.argsort(-scores)
313
+ keep: list[int] = []
314
+ while len(order) > 0:
315
+ i = int(order[0])
316
+ keep.append(i)
317
+ if len(order) == 1:
318
+ break
319
+ rest = order[1:]
320
+ xx1 = np.maximum(boxes[i, 0], boxes[rest, 0])
321
+ yy1 = np.maximum(boxes[i, 1], boxes[rest, 1])
322
+ xx2 = np.minimum(boxes[i, 2], boxes[rest, 2])
323
+ yy2 = np.minimum(boxes[i, 3], boxes[rest, 3])
324
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
325
+ a_i = (max(0.0, boxes[i, 2] - boxes[i, 0]) *
326
+ max(0.0, boxes[i, 3] - boxes[i, 1]))
327
+ a_r = (np.maximum(0.0, boxes[rest, 2] - boxes[rest, 0]) *
328
+ np.maximum(0.0, boxes[rest, 3] - boxes[rest, 1]))
329
+ iou = inter / (a_i + a_r - inter + 1e-7)
330
+ order = rest[iou <= iou_thresh]
331
+ return np.array(keep, dtype=np.intp)
332
+
333
+ def _per_class_hard_nms(
334
+ self,
335
+ boxes: np.ndarray,
336
+ scores: np.ndarray,
337
+ cls_ids: np.ndarray,
338
+ iou_thresh: float,
339
+ ) -> np.ndarray:
340
+ if len(boxes) == 0:
341
+ return np.array([], dtype=np.intp)
342
+ all_keep: list[int] = []
343
+ for c in np.unique(cls_ids):
344
+ mask = cls_ids == c
345
+ indices = np.where(mask)[0]
346
+ keep = self._hard_nms(boxes[mask], scores[mask], iou_thresh)
347
+ all_keep.extend(indices[keep].tolist())
348
+ all_keep.sort()
349
+ return np.array(all_keep, dtype=np.intp)
350
+
351
+ def _soft_nms(
352
+ self,
353
+ boxes: np.ndarray,
354
+ scores: np.ndarray,
355
+ sigma: float = 0.5,
356
+ score_thresh: float = 0.01,
357
+ ) -> tuple[np.ndarray, np.ndarray]:
358
+ """Soft-NMS: Gaussian decay of overlapping scores instead of hard removal.
359
+ Returns (kept_original_indices, updated_scores). (Ported from carwash001.)"""
360
+ N = len(boxes)
361
+ if N == 0:
362
+ return np.array([], dtype=np.intp), np.array([], dtype=np.float32)
363
+ boxes = boxes.astype(np.float32, copy=True)
364
+ scores = scores.astype(np.float32, copy=True)
365
+ order = np.arange(N)
366
+ for i in range(N):
367
+ max_pos = i + int(np.argmax(scores[i:]))
368
+ boxes[[i, max_pos]] = boxes[[max_pos, i]]
369
+ scores[[i, max_pos]] = scores[[max_pos, i]]
370
+ order[[i, max_pos]] = order[[max_pos, i]]
371
+ if i + 1 >= N:
372
+ break
373
+ xx1 = np.maximum(boxes[i, 0], boxes[i + 1:, 0])
374
+ yy1 = np.maximum(boxes[i, 1], boxes[i + 1:, 1])
375
+ xx2 = np.minimum(boxes[i, 2], boxes[i + 1:, 2])
376
+ yy2 = np.minimum(boxes[i, 3], boxes[i + 1:, 3])
377
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
378
+ area_i = max(0.0, float(
379
+ (boxes[i, 2] - boxes[i, 0]) * (boxes[i, 3] - boxes[i, 1])))
380
+ areas_j = (np.maximum(0.0, boxes[i + 1:, 2] - boxes[i + 1:, 0])
381
+ * np.maximum(0.0, boxes[i + 1:, 3] - boxes[i + 1:, 1]))
382
+ iou = inter / (area_i + areas_j - inter + 1e-7)
383
+ scores[i + 1:] *= np.exp(-(iou ** 2) / sigma)
384
+ mask = scores > score_thresh
385
+ return order[mask], scores[mask]
386
+
387
+ def _per_class_soft_nms(
388
+ self,
389
+ boxes: np.ndarray,
390
+ scores: np.ndarray,
391
+ cls_ids: np.ndarray,
392
+ sigma: float = 0.5,
393
+ score_thresh: float = 0.01,
394
+ ) -> tuple[np.ndarray, np.ndarray]:
395
+ """Soft-NMS applied independently per class. Returns (kept_idx, updated_scores)."""
396
+ if len(boxes) == 0:
397
+ return np.array([], dtype=np.intp), np.array([], dtype=np.float32)
398
+ all_keep: list[int] = []
399
+ all_scores: list[float] = []
400
+ for c in np.unique(cls_ids):
401
+ indices = np.where(cls_ids == c)[0]
402
+ keep, updated = self._soft_nms(boxes[indices], scores[indices],
403
+ sigma, score_thresh)
404
+ for k, s in zip(keep, updated):
405
+ all_keep.append(int(indices[k])); all_scores.append(float(s))
406
+ if not all_keep:
407
+ return np.array([], dtype=np.intp), np.array([], dtype=np.float32)
408
+ return np.array(all_keep, dtype=np.intp), np.array(all_scores, dtype=np.float32)
409
+
410
+ def _cross_class_dedup_op(
411
+ self,
412
+ boxes: np.ndarray,
413
+ scores: np.ndarray,
414
+ cls_ids: np.ndarray,
415
+ iou_thresh: float,
416
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
417
+ """Remove near-duplicate boxes across classes.
418
+
419
+ Order candidates by (score - per_class_threshold) margin, then by area;
420
+ keep the highest, suppress every other box with IoU > iou_thresh.
421
+ With a single road_sign class this is effectively a no-op, but the
422
+ method is kept so the pipeline stays compatible with the multi-class
423
+ miner template.
424
+ """
425
+ n = len(boxes)
426
+ if n <= 1:
427
+ return boxes, scores, cls_ids
428
+ boxes = np.asarray(boxes, dtype=np.float32)
429
+ scores = np.asarray(scores, dtype=np.float32)
430
+ cls_ids = np.asarray(cls_ids, dtype=np.int32)
431
+ areas = (np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) *
432
+ np.maximum(0.0, boxes[:, 3] - boxes[:, 1]))
433
+ margins = scores - self._conf_thres_array[cls_ids]
434
+ order = np.lexsort((-areas, -margins))
435
+ suppressed = np.zeros(n, dtype=bool)
436
+ keep: list[int] = []
437
+ for i in order:
438
+ if suppressed[i]:
439
+ continue
440
+ keep.append(int(i))
441
+ bi = boxes[i]
442
+ xx1 = np.maximum(bi[0], boxes[:, 0])
443
+ yy1 = np.maximum(bi[1], boxes[:, 1])
444
+ xx2 = np.minimum(bi[2], boxes[:, 2])
445
+ yy2 = np.minimum(bi[3], boxes[:, 3])
446
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
447
+ a_i = max(1e-7, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
448
+ iou = inter / (a_i + areas - inter + 1e-7)
449
+ dup = iou > iou_thresh
450
+ dup[i] = False
451
+ suppressed |= dup
452
+ keep_idx = np.array(keep, dtype=np.intp)
453
+ return boxes[keep_idx], scores[keep_idx], cls_ids[keep_idx]
454
+
455
+ @staticmethod
456
+ def _max_score_per_cluster(
457
+ post_boxes: np.ndarray,
458
+ post_cls: np.ndarray,
459
+ full_boxes: np.ndarray,
460
+ full_scores: np.ndarray,
461
+ full_cls: np.ndarray,
462
+ iou_thresh: float,
463
+ ) -> np.ndarray:
464
+ """For each kept (post-NMS) box, return the max score over the FULL
465
+ candidate set among same-class boxes with IoU >= iou_thresh.
466
+
467
+ Used after horizontal-flip TTA: a high-confidence flipped detection
468
+ can raise the score of the corresponding original detection.
469
+ """
470
+ n = len(post_boxes)
471
+ if n == 0:
472
+ return np.empty(0, dtype=np.float32)
473
+ full_areas = (np.maximum(0.0, full_boxes[:, 2] - full_boxes[:, 0]) *
474
+ np.maximum(0.0, full_boxes[:, 3] - full_boxes[:, 1]))
475
+ out = np.empty(n, dtype=np.float32)
476
+ for i in range(n):
477
+ bi = post_boxes[i]
478
+ xx1 = np.maximum(bi[0], full_boxes[:, 0])
479
+ yy1 = np.maximum(bi[1], full_boxes[:, 1])
480
+ xx2 = np.minimum(bi[2], full_boxes[:, 2])
481
+ yy2 = np.minimum(bi[3], full_boxes[:, 3])
482
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
483
+ a_i = max(0.0, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
484
+ iou = inter / (a_i + full_areas - inter + 1e-7)
485
+ cluster = (iou >= iou_thresh) & (full_cls == post_cls[i])
486
+ out[i] = float(np.max(full_scores[cluster])) if np.any(cluster) else 0.0
487
+ return out
488
+
489
+ def _conf_filter_mask(
490
+ self, scores: np.ndarray, cls_ids: np.ndarray
491
+ ) -> np.ndarray:
492
+ """Boolean keep-mask: score >= per-class threshold, with a per-class
493
+ rescue -- if a class has zero boxes passing, admit its top-1 candidate
494
+ when its score >= (per-class threshold - per-class bonus)."""
495
+ if len(scores) == 0:
496
+ return np.zeros(0, dtype=bool)
497
+ thr = self._conf_thres_array[cls_ids]
498
+ keep = scores >= thr
499
+ for c in np.unique(cls_ids):
500
+ b = float(self._bonus_array[c])
501
+ if b <= 0.0:
502
+ continue
503
+ cm = cls_ids == c
504
+ if keep[cm].any():
505
+ continue
506
+ idx = np.where(cm)[0]
507
+ top = int(idx[int(np.argmax(scores[idx]))])
508
+ if scores[top] >= self._conf_thres_array[c] - b:
509
+ keep[top] = True
510
+ return keep
511
+
512
+ def _filter_sane_boxes(
513
+ self,
514
+ boxes: np.ndarray,
515
+ scores: np.ndarray,
516
+ cls_ids: np.ndarray,
517
+ orig_size: tuple[int, int],
518
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
519
+ """Drop tiny / degenerate / image-spanning / extreme-AR boxes (FP)."""
520
+ if len(boxes) == 0:
521
+ return boxes, scores, cls_ids
522
+ orig_w, orig_h = orig_size
523
+ image_area = float(orig_w * orig_h)
524
+ keep = []
525
+ for i, box in enumerate(boxes):
526
+ x1, y1, x2, y2 = box.tolist()
527
+ bw = x2 - x1
528
+ bh = y2 - y1
529
+ if bw <= 0 or bh <= 0:
530
+ continue
531
+ if bw < self.min_side or bh < self.min_side:
532
+ continue
533
+ area = bw * bh
534
+ if area < self.min_box_area:
535
+ continue
536
+ if area > 0.95 * image_area:
537
+ continue
538
+ ar = max(bw / max(bh, 1e-6), bh / max(bw, 1e-6))
539
+ if ar > self.max_aspect_ratio:
540
+ continue
541
+ keep.append(i)
542
+ if not keep:
543
+ return (
544
+ np.empty((0, 4), dtype=np.float32),
545
+ np.empty((0,), dtype=np.float32),
546
+ np.empty((0,), dtype=np.int32),
547
+ )
548
+ k = np.array(keep, dtype=np.intp)
549
+ return boxes[k], scores[k], cls_ids[k]
550
+
551
+ def _per_view_pipeline(
552
+ self,
553
+ boxes: np.ndarray,
554
+ scores: np.ndarray,
555
+ cls_ids: np.ndarray,
556
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
557
+ """Per-view post-processing pipeline: per-class NMS -> cap -> cross-class dedup."""
558
+ if len(boxes) > 1:
559
+ if self.use_soft_nms:
560
+ keep, new_scores = self._per_class_soft_nms(
561
+ boxes, scores, cls_ids,
562
+ self.soft_nms_sigma, self.soft_nms_score_thresh)
563
+ boxes, scores, cls_ids = boxes[keep], new_scores, cls_ids[keep]
564
+ else:
565
+ keep = self._per_class_hard_nms(boxes, scores, cls_ids, self.iou_thres)
566
+ boxes, scores, cls_ids = boxes[keep], scores[keep], cls_ids[keep]
567
+ if len(scores) > self.max_det:
568
+ top = np.argsort(-scores)[: self.max_det]
569
+ boxes, scores, cls_ids = boxes[top], scores[top], cls_ids[top]
570
+ if len(boxes) > 1:
571
+ boxes, scores, cls_ids = self._cross_class_dedup_op(
572
+ boxes, scores, cls_ids, self.cross_iou_thresh
573
+ )
574
+ return boxes, scores, cls_ids
575
+
576
+ @staticmethod
577
+ def _build_results(
578
+ boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray
579
+ ) -> list[BoundingBox]:
580
+ results: list[BoundingBox] = []
581
+ for box, conf, cls_id in zip(boxes, scores, cls_ids):
582
+ x1, y1, x2, y2 = box.tolist()
583
+ if x2 <= x1 or y2 <= y1:
584
+ continue
585
+ results.append(
586
+ BoundingBox(
587
+ x1=int(math.floor(x1)),
588
+ y1=int(math.floor(y1)),
589
+ x2=int(math.ceil(x2)),
590
+ y2=int(math.ceil(y2)),
591
+ cls_id=int(cls_id),
592
+ conf=float(conf),
593
+ )
594
+ )
595
+ return results
596
+
597
+ def _decode_final_dets(
598
+ self,
599
+ preds: np.ndarray,
600
+ ratio: float,
601
+ pad: tuple[float, float],
602
+ orig_size: tuple[int, int],
603
+ ) -> list[BoundingBox]:
604
+
605
+ # if self.use_fast_postprocess:
606
+ # return self._decode_final_dets_fast(preds, ratio, pad, orig_size)
607
+
608
+ """Final-detection output path: rows shaped [x1, y1, x2, y2, conf, cls_id]."""
609
+ if preds.ndim == 3 and preds.shape[0] == 1:
610
+ preds = preds[0]
611
+ if preds.ndim != 2 or preds.shape[1] < 6:
612
+ raise ValueError(f"Unexpected ONNX final-det output shape: {preds.shape}")
613
+
614
+ boxes = preds[:, :4].astype(np.float32)
615
+ scores = preds[:, 4].astype(np.float32)
616
+ cls_ids = preds[:, 5].astype(np.int32)
617
+ cls_ids = self.cls_remap[cls_ids]
618
+
619
+ keep = self._conf_filter_mask(scores, cls_ids)
620
+ boxes = boxes[keep]
621
+ scores = scores[keep]
622
+ cls_ids = cls_ids[keep]
623
+ if len(boxes) == 0:
624
+ return []
625
+
626
+ pad_w, pad_h = pad
627
+ boxes[:, [0, 2]] -= pad_w
628
+ boxes[:, [1, 3]] -= pad_h
629
+ boxes /= ratio
630
+ boxes = self._clip_boxes(boxes, orig_size)
631
+
632
+ boxes, scores, cls_ids = self._filter_sane_boxes(
633
+ boxes, scores, cls_ids, orig_size
634
+ )
635
+ if len(boxes) == 0:
636
+ return []
637
+ boxes, scores, cls_ids = self._per_view_pipeline(boxes, scores, cls_ids)
638
+ return self._build_results(boxes, scores, cls_ids)
639
+
640
+ def _decode_raw_yolo(
641
+ self,
642
+ preds: np.ndarray,
643
+ ratio: float,
644
+ pad: tuple[float, float],
645
+ orig_size: tuple[int, int],
646
+ ) -> list[BoundingBox]:
647
+ """Fallback raw-YOLO output path: per-anchor class logits."""
648
+ if preds.ndim != 3 or preds.shape[0] != 1:
649
+ raise ValueError(f"Unexpected raw ONNX output shape: {preds.shape}")
650
+ preds = preds[0]
651
+ if preds.shape[0] <= 16 and preds.shape[1] > preds.shape[0]:
652
+ preds = preds.T
653
+ if preds.ndim != 2 or preds.shape[1] < 5:
654
+ raise ValueError(f"Unexpected raw output shape: {preds.shape}")
655
+
656
+ boxes_xywh = preds[:, :4].astype(np.float32)
657
+ cls_part = preds[:, 4:].astype(np.float32)
658
+ if cls_part.shape[1] == 1:
659
+ scores = cls_part[:, 0]
660
+ cls_ids = np.zeros(len(scores), dtype=np.int32)
661
+ else:
662
+ cls_ids = np.argmax(cls_part, axis=1).astype(np.int32)
663
+ scores = cls_part[np.arange(len(cls_part)), cls_ids]
664
+ cls_ids = self.cls_remap[cls_ids]
665
+
666
+ keep = self._conf_filter_mask(scores, cls_ids)
667
+ boxes_xywh = boxes_xywh[keep]
668
+ scores = scores[keep]
669
+ cls_ids = cls_ids[keep]
670
+ if len(boxes_xywh) == 0:
671
+ return []
672
+ boxes = self._xywh_to_xyxy(boxes_xywh)
673
+
674
+ pad_w, pad_h = pad
675
+ boxes[:, [0, 2]] -= pad_w
676
+ boxes[:, [1, 3]] -= pad_h
677
+ boxes /= ratio
678
+ boxes = self._clip_boxes(boxes, orig_size)
679
+
680
+ boxes, scores, cls_ids = self._filter_sane_boxes(
681
+ boxes, scores, cls_ids, orig_size
682
+ )
683
+ if len(boxes) == 0:
684
+ return []
685
+
686
+ boxes, scores, cls_ids = self._per_view_pipeline(boxes, scores, cls_ids)
687
+ return self._build_results(boxes, scores, cls_ids)
688
+
689
+ def _postprocess(
690
+ self,
691
+ output: np.ndarray,
692
+ ratio: float,
693
+ pad: tuple[float, float],
694
+ orig_size: tuple[int, int],
695
+ ) -> list[BoundingBox]:
696
+ if output.ndim == 2 and output.shape[1] >= 6:
697
+ return self._decode_final_dets(output, ratio, pad, orig_size)
698
+ if output.ndim == 3 and output.shape[0] == 1 and output.shape[2] == 6:
699
+ return self._decode_final_dets(output, ratio, pad, orig_size)
700
+ return self._decode_raw_yolo(output, ratio, pad, orig_size)
701
+
702
+ @staticmethod
703
+ def _iou_box(a: BoundingBox, b: BoundingBox) -> float:
704
+ x1 = max(a.x1, b.x1)
705
+ y1 = max(a.y1, b.y1)
706
+ x2 = min(a.x2, b.x2)
707
+ y2 = min(a.y2, b.y2)
708
+ inter = max(0, x2 - x1) * max(0, y2 - y1)
709
+ if inter <= 0:
710
+ return 0.0
711
+ area_a = max(0, a.x2 - a.x1) * max(0, a.y2 - a.y1)
712
+ area_b = max(0, b.x2 - b.x1) * max(0, b.y2 - b.y1)
713
+ union = area_a + area_b - inter
714
+ return float(inter / union) if union > 0 else 0.0
715
+
716
+ def _selective_merge(
717
+ self,
718
+ primary: list[BoundingBox],
719
+ secondary_boxes: list[BoundingBox],
720
+ ) -> list[BoundingBox]:
721
+ """Keep high-conf secondary boxes that do not overlap primary."""
722
+ if not secondary_boxes:
723
+ return []
724
+ merged: list[BoundingBox] = []
725
+ for a in secondary_boxes:
726
+ if a.conf < self.secondary_conf:
727
+ continue
728
+ if primary:
729
+ max_iou = max(self._iou_box(a, p) for p in primary)
730
+ if max_iou >= self.merge_iou:
731
+ continue
732
+ merged.append(a)
733
+ return merged
734
+
735
+ def _harmonize_overlapping_heads(
736
+ self,
737
+ primary: list[BoundingBox],
738
+ secondary: list[BoundingBox],
739
+ ) -> list[BoundingBox]:
740
+ """Where primary and secondary overlap, keep primary geometry and partner score."""
741
+ if not primary or not secondary:
742
+ return primary
743
+
744
+ out: list[BoundingBox] = []
745
+ for p in primary:
746
+ partner_conf = p.conf
747
+ best_iou = 0.0
748
+ for s in secondary:
749
+ if s.cls_id != p.cls_id:
750
+ continue
751
+ iou = self._iou_box(p, s)
752
+ if iou >= self.merge_iou and iou > best_iou:
753
+ best_iou = iou
754
+ partner_conf = s.conf if s.conf >= 0.90 else p.conf
755
+ out.append(
756
+ BoundingBox(
757
+ x1=p.x1,
758
+ y1=p.y1,
759
+ x2=p.x2,
760
+ y2=p.y2,
761
+ cls_id=p.cls_id,
762
+ conf=float(partner_conf),
763
+ )
764
+ )
765
+ return out
766
+
767
+ def _merge_dual_head_boxes(
768
+ self,
769
+ primary_boxes: list[BoundingBox],
770
+ secondary_boxes: list[BoundingBox],
771
+ orig_size: tuple[int, int],
772
+ ) -> list[BoundingBox]:
773
+ """Preserve primary boxes; append alex fill after extra-only dedupe.
774
+
775
+ `merge_iou` only gates cross-head exclusion in `_selective_merge`.
776
+ `dual_head_dedup_iou` only dedupes alex extras (never re-NMS primary).
777
+ Primary geometry/conf are left unchanged (no harmonize).
778
+ """
779
+ extra = self._selective_merge(primary_boxes, secondary_boxes)
780
+ if not extra:
781
+ return primary_boxes
782
+
783
+ deduped_extra = self._merge_views(
784
+ [extra], orig_size, iou_thresh=self.dual_head_dedup_iou
785
+ )
786
+ return primary_boxes + deduped_extra
787
+
788
+ def _merge_model_outputs(
789
+ self,
790
+ primary_out: np.ndarray,
791
+ secondary_out: np.ndarray,
792
+ ratio: float,
793
+ pad: tuple[float, float],
794
+ orig_size: tuple[int, int],
795
+ ) -> list[BoundingBox]:
796
+ """Primary detections plus optional non-overlapping secondary boxes."""
797
+ primary_boxes = self._postprocess(primary_out, ratio, pad, orig_size)
798
+ if not self.use_secondary_merge:
799
+ return primary_boxes
800
+ secondary_boxes = self._postprocess(secondary_out, ratio, pad, orig_size)
801
+ return self._merge_dual_head_boxes(
802
+ primary_boxes, secondary_boxes, orig_size
803
+ )
804
+
805
+ def _predict_single(self, image: np.ndarray) -> list[BoundingBox]:
806
+ if image is None:
807
+ raise ValueError("Input image is None")
808
+ if not isinstance(image, np.ndarray):
809
+ raise TypeError(f"Input is not numpy array: {type(image)}")
810
+ if image.ndim != 3:
811
+ raise ValueError(f"Expected HWC image, got shape={image.shape}")
812
+ if image.shape[0] <= 0 or image.shape[1] <= 0:
813
+ raise ValueError(f"Invalid image shape={image.shape}")
814
+ if image.shape[2] != 3:
815
+ raise ValueError(f"Expected 3 channels, got shape={image.shape}")
816
+ if image.dtype != np.uint8:
817
+ image = image.astype(np.uint8)
818
+
819
+ input_tensor, ratio, pad, orig_size = self._preprocess(image)
820
+ expected = (1, 3, self.input_height, self.input_width)
821
+ if input_tensor.shape != expected:
822
+ raise ValueError(
823
+ f"Bad input tensor shape={input_tensor.shape}, expected={expected}"
824
+ )
825
+
826
+ outputs = self.session.run(self.output_names, {self.input_name: input_tensor})
827
+ if len(outputs) < 2:
828
+ return self._postprocess(outputs[0], ratio, pad, orig_size)
829
+
830
+ out0, out1 = outputs[0], outputs[1]
831
+ # Drop batch dim when present ([1, N, 6] -> [N, 6]).
832
+ if isinstance(out0, np.ndarray) and out0.ndim == 3 and out0.shape[0] == 1:
833
+ out0 = out0[0]
834
+ if isinstance(out1, np.ndarray) and out1.ndim == 3 and out1.shape[0] == 1:
835
+ out1 = out1[0]
836
+ return self._merge_model_outputs(out0, out1, ratio, pad, orig_size)
837
+
838
+ def _predict_tta(self, image: np.ndarray) -> list[BoundingBox]:
839
+ """Horizontal-flip TTA.
840
+
841
+ Strategy:
842
+ 1. Predict on original and on flipped image.
843
+ 2. Map flipped boxes back to original coordinates.
844
+ 3. Per-class hard NMS on the union.
845
+ 4. For each kept box, compute the max same-class score across the
846
+ FULL union (not just the post-NMS subset) -- this lets a high-
847
+ confidence flipped detection raise a borderline original one.
848
+ 5. Cross-class dedup to suppress same-physical-object multi-class.
849
+ """
850
+ boxes_orig = self._predict_single(image)
851
+ flipped = cv2.flip(image, 1)
852
+ boxes_flip = self._predict_single(flipped)
853
+ w = image.shape[1]
854
+ boxes_flip = [
855
+ BoundingBox(
856
+ x1=w - b.x2, y1=b.y1, x2=w - b.x1, y2=b.y2,
857
+ cls_id=b.cls_id, conf=b.conf,
858
+ )
859
+ for b in boxes_flip
860
+ ]
861
+ all_boxes = boxes_orig + boxes_flip
862
+ if not all_boxes:
863
+ return []
864
+
865
+ coords = np.array(
866
+ [[b.x1, b.y1, b.x2, b.y2] for b in all_boxes], dtype=np.float32
867
+ )
868
+ scores = np.array([b.conf for b in all_boxes], dtype=np.float32)
869
+ cls_ids = np.array([b.cls_id for b in all_boxes], dtype=np.int32)
870
+
871
+ hard_keep = self._per_class_hard_nms(coords, scores, cls_ids, self.iou_thres)
872
+ if len(hard_keep) == 0:
873
+ return []
874
+ if len(hard_keep) > self.max_det:
875
+ top = np.argsort(-scores[hard_keep])[: self.max_det]
876
+ hard_keep = hard_keep[top]
877
+
878
+ boosted = self._max_score_per_cluster(
879
+ coords[hard_keep], cls_ids[hard_keep],
880
+ coords, scores, cls_ids, self.iou_thres,
881
+ )
882
+
883
+ kept_coords = coords[hard_keep]
884
+ kept_cls = cls_ids[hard_keep]
885
+ if len(kept_coords) > 1:
886
+ kept_coords, boosted, kept_cls = self._cross_class_dedup_op(
887
+ kept_coords, boosted, kept_cls, self.cross_iou_thresh
888
+ )
889
+
890
+ return [
891
+ BoundingBox(
892
+ x1=int(math.floor(kept_coords[j, 0])),
893
+ y1=int(math.floor(kept_coords[j, 1])),
894
+ x2=int(math.ceil(kept_coords[j, 2])),
895
+ y2=int(math.ceil(kept_coords[j, 3])),
896
+ cls_id=int(kept_cls[j]),
897
+ conf=float(boosted[j]),
898
+ )
899
+ for j in range(len(kept_coords))
900
+ ]
901
+
902
+ def _predict_tiles(self, image: np.ndarray) -> list[BoundingBox]:
903
+ """Tile-based TTA for high-resolution images.
904
+
905
+ Splits the source image into two overlapping horizontal tiles, runs
906
+ single-pass inference on each at native scale, and translates boxes
907
+ back to the global frame. Useful when source width >> model input
908
+ width because letterboxing otherwise discards effective resolution
909
+ that small / distant signs depend on.
910
+
911
+ Returns an empty list if the image isn't wide enough to benefit; the
912
+ caller falls back to the regular pipeline in that case.
913
+ """
914
+ h, w = image.shape[:2]
915
+ if w < int(self.input_width * self.tile_trigger_ratio):
916
+ return []
917
+
918
+ overlap = int(w * self.tile_overlap_ratio)
919
+ mid = w // 2
920
+ x_left_end = min(w, mid + overlap // 2)
921
+ x_right_start = max(0, mid - overlap // 2)
922
+
923
+ left = image[:, :x_left_end]
924
+ right = image[:, x_right_start:]
925
+
926
+ boxes_left = self._predict_single(left)
927
+ boxes_right = self._predict_single(right)
928
+
929
+ shifted_right = [
930
+ BoundingBox(
931
+ x1=b.x1 + x_right_start,
932
+ y1=b.y1,
933
+ x2=b.x2 + x_right_start,
934
+ y2=b.y2,
935
+ cls_id=b.cls_id,
936
+ conf=b.conf,
937
+ )
938
+ for b in boxes_right
939
+ ]
940
+ return boxes_left + shifted_right
941
+
942
+ def _merge_views(
943
+ self,
944
+ view_boxes: list[list[BoundingBox]],
945
+ image_size: tuple[int, int],
946
+ iou_thresh: float | None = None,
947
+ ) -> list[BoundingBox]:
948
+ """Merge boxes from multiple views (single / hflip / tiles).
949
+
950
+ Same logic as `_predict_tta`'s tail: per-class hard NMS to dedupe,
951
+ then for each kept box take the max same-class score across the full
952
+ candidate union — a high-confidence detection in any view boosts
953
+ borderline matches in others.
954
+ """
955
+ nms_iou = self.iou_thres if iou_thresh is None else float(iou_thresh)
956
+ all_boxes: list[BoundingBox] = []
957
+ for vb in view_boxes:
958
+ all_boxes.extend(vb)
959
+ if not all_boxes:
960
+ return []
961
+
962
+ coords = np.array(
963
+ [[b.x1, b.y1, b.x2, b.y2] for b in all_boxes], dtype=np.float32
964
+ )
965
+ scores = np.array([b.conf for b in all_boxes], dtype=np.float32)
966
+ cls_ids = np.array([b.cls_id for b in all_boxes], dtype=np.int32)
967
+
968
+ coords = self._clip_boxes(coords, image_size)
969
+
970
+ hard_keep = self._per_class_hard_nms(coords, scores, cls_ids, nms_iou)
971
+ if len(hard_keep) == 0:
972
+ return []
973
+ if len(hard_keep) > self.max_det:
974
+ top = np.argsort(-scores[hard_keep])[: self.max_det]
975
+ hard_keep = hard_keep[top]
976
+
977
+ boosted = self._max_score_per_cluster(
978
+ coords[hard_keep], cls_ids[hard_keep],
979
+ coords, scores, cls_ids, nms_iou,
980
+ )
981
+
982
+ kept_coords = coords[hard_keep]
983
+ kept_cls = cls_ids[hard_keep]
984
+ if len(kept_coords) > 1:
985
+ kept_coords, boosted, kept_cls = self._cross_class_dedup_op(
986
+ kept_coords, boosted, kept_cls, self.cross_iou_thresh
987
+ )
988
+
989
+ return [
990
+ BoundingBox(
991
+ x1=int(math.floor(kept_coords[j, 0])),
992
+ y1=int(math.floor(kept_coords[j, 1])),
993
+ x2=int(math.ceil(kept_coords[j, 2])),
994
+ y2=int(math.ceil(kept_coords[j, 3])),
995
+ cls_id=int(kept_cls[j]),
996
+ conf=float(boosted[j]),
997
+ )
998
+ for j in range(len(kept_coords))
999
+ ]
1000
+
1001
+ def _predict_full(self, image: np.ndarray) -> list[BoundingBox]:
1002
+ """Top-level per-frame prediction with all enabled augmentations.
1003
+
1004
+ - `use_tta=True`: original + horizontal flip
1005
+ - `use_tile_tta=True` AND image wide enough: two overlapping tiles
1006
+ All views are merged via per-class NMS + cluster-max score boost.
1007
+ """
1008
+ h, w = image.shape[:2]
1009
+ image_size = (w, h)
1010
+
1011
+ if not self.use_tta and not self.use_tile_tta:
1012
+ return self._finalize_boxes(self._predict_single(image), image_size)
1013
+
1014
+ views: list[list[BoundingBox]] = []
1015
+ if self.use_tta:
1016
+ views.append(self._predict_single(image))
1017
+ flipped = cv2.flip(image, 1)
1018
+ flipped_dets = self._predict_single(flipped)
1019
+ views.append([
1020
+ BoundingBox(
1021
+ x1=w - b.x2, y1=b.y1, x2=w - b.x1, y2=b.y2,
1022
+ cls_id=b.cls_id, conf=b.conf,
1023
+ )
1024
+ for b in flipped_dets
1025
+ ])
1026
+ else:
1027
+ views.append(self._predict_single(image))
1028
+
1029
+ if self.use_tile_tta:
1030
+ tile_boxes = self._predict_tiles(image)
1031
+ if tile_boxes:
1032
+ views.append(tile_boxes)
1033
+
1034
+ return self._finalize_boxes(self._merge_views(views, image_size), image_size)
1035
+
1036
+ def predict_batch(
1037
+ self,
1038
+ batch_images: list[ndarray],
1039
+ offset: int,
1040
+ n_keypoints: int,
1041
+ ) -> list[TVFrameResult]:
1042
+ results: list[TVFrameResult] = []
1043
+ for frame_number_in_batch, image in enumerate(batch_images):
1044
+ try:
1045
+ boxes = self._predict_full(image)
1046
+ except Exception as e:
1047
+ print(
1048
+ f"⚠️ Inference failed for frame "
1049
+ f"{offset + frame_number_in_batch}: {e}"
1050
+ )
1051
+ boxes = []
1052
+ results.append(
1053
+ TVFrameResult(
1054
+ frame_id=offset + frame_number_in_batch,
1055
+ boxes=boxes,
1056
+ keypoints=[(0, 0) for _ in range(max(0, int(n_keypoints)))],
1057
+ )
1058
+ )
1059
+ return results