SuperBitDev commited on
Commit
2718280
·
verified ·
1 Parent(s): 96cabc4

Upload folder using huggingface_hub

Browse files
Files changed (3) hide show
  1. chute_config.yml +1 -0
  2. miner.py +911 -937
  3. weights.onnx +2 -2
chute_config.yml CHANGED
@@ -8,6 +8,7 @@ Image:
8
  NodeSelector:
9
  gpu_count: 1
10
  min_vram_gb_per_gpu: 16
 
11
  include:
12
  - pro_6000
13
 
 
8
  NodeSelector:
9
  gpu_count: 1
10
  min_vram_gb_per_gpu: 16
11
+ max_hourly_price_per_gpu: 2
12
  include:
13
  - pro_6000
14
 
miner.py CHANGED
@@ -1,937 +1,911 @@
1
- from pathlib import Path
2
-
3
- import cv2
4
- import numpy as np
5
- import onnxruntime as ort
6
- from numpy import ndarray
7
- from pydantic import BaseModel
8
-
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
- Strategy (ported from offense / fire001 miner):
28
- - per-class confidence threshold with per-class rescue bonus
29
- - per-class hard NMS, then cross-class dedup (no-op for single class)
30
- - horizontal-flip TTA with full-set cluster score boost
31
- Plus: class remap, sanity-box filter tuned for small distant signs,
32
- TTA toggle.
33
- """
34
-
35
- class_names = ["road_sign"]
36
- # Order the model emits classes in -- remapped to `class_names` index.
37
- _model_class_order = ["road_sign"]
38
-
39
- iou_thres = 0.5
40
- cross_iou_thresh = 0.8
41
- max_det = 150
42
-
43
- # Per-class confidence threshold. Road signs in this dataset are
44
- # frequently degraded / rear-facing / partly-obscured / distant, so we
45
- # run noticeably below the fire/smoke baseline. The validator's
46
- # false_positive pillar = max(0, 1 - ffpi/10): we can tolerate ~2 FP per
47
- # image and still keep that pillar above 0.8.
48
- _conf_thres_array = np.array(
49
- [0.33], dtype=np.float32
50
- )
51
- # Per-class rescue bonus. If a class has ZERO boxes passing the threshold
52
- # in a frame, its top-1 candidate is admitted when its score is at least
53
- # (threshold - bonus). Bumped from 0.05 -> 0.10 so a single faint sign in
54
- # an otherwise empty frame still produces a detection (map50 recall win,
55
- # at most one extra FP per such frame).
56
- _bonus_array = np.array(
57
- [0.12], dtype=np.float32
58
- )
59
-
60
- # Box sanity filter: drop tiny / degenerate / image-spanning / extreme
61
- # aspect ratio boxes.
62
- # min_box_area = 14x14 -> 14x14 is the smallest credible sign. The old
63
- # value of 64 (8x8) silently discarded narrow
64
- # distant signs like a 10x6 px overhead chevron.
65
- # min_side = 3 -> matches min_box_area; anything thinner is
66
- # almost certainly a pole or shadow false alarm.
67
- # max_aspect_ratio = 12.0
68
- # -> overhead destination panels and lane-assignment
69
- # signs are very wide (long, thin rectangles);
70
- # 8.0 was clipping legitimate detections.
71
- min_box_area = 8 * 8
72
- min_side = 3
73
- max_aspect_ratio = 12.0
74
-
75
- # Final box-size calibration. The detector + de-letterbox + integer-rounding
76
- # pipeline emits boxes slightly larger than the object, so shrink every
77
- # emitted box about its center by a fixed per-axis factor before output:
78
- # new_w = w / box_shrink_w, new_h = h / box_shrink_h.
79
- box_shrink_w = 1.027
80
- box_shrink_h = 1.014
81
-
82
- # Tile-based TTA: when the source image is significantly larger than the
83
- # model input, letterboxing throws away ~1.5x of effective resolution,
84
- # which kills small-sign recall. Splitting into overlapping horizontal
85
- # tiles preserves native resolution on each half. Triggered only when
86
- # source width >= tile_trigger_ratio * model_input_width to avoid wasted
87
- # compute on already-small images.
88
- tile_trigger_ratio = 1.4
89
- tile_overlap_ratio = 0.20
90
-
91
- def __init__(self, path_hf_repo: Path) -> None:
92
- model_path = path_hf_repo / "weights.onnx"
93
-
94
- print("ORT version:", ort.__version__)
95
-
96
- try:
97
- ort.preload_dlls()
98
- print("✅ onnxruntime.preload_dlls() success")
99
- except Exception as e:
100
- print(f"⚠️ preload_dlls failed: {e}")
101
-
102
- print("ORT available providers BEFORE session:", ort.get_available_providers())
103
-
104
- sess_options = ort.SessionOptions()
105
- sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
106
- sess_options.intra_op_num_threads = 2
107
- sess_options.inter_op_num_threads = 1
108
- sess_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
109
-
110
- try:
111
- self.session = ort.InferenceSession(
112
- str(model_path),
113
- sess_options=sess_options,
114
- providers=["CPUExecutionProvider"],
115
- )
116
- except Exception as e:
117
- self.session = ort.InferenceSession(
118
- str(model_path),
119
- sess_options=sess_options,
120
- providers=["CPUExecutionProvider"],
121
- )
122
-
123
- print("ORT session providers:", self.session.get_providers())
124
-
125
- # Build cls_remap: for each model-emit index i,
126
- # cls_remap[i] = self.class_names.index(model_class_order[i])
127
- # i.e. convert a model-side class id into the output class id that
128
- # downstream code (BoundingBox.cls_id, the per-class threshold/bonus
129
- # arrays) expects. The model-side order comes from the ONNX metadata
130
- # when available, else falls back to the static _model_class_order.
131
- model_class_order = self._read_model_class_order()
132
- if model_class_order is None:
133
- model_class_order = list(self._model_class_order)
134
- print(f"cls order: no usable ONNX metadata, FALLBACK {model_class_order}")
135
- else:
136
- print(f"cls order: from ONNX metadata {model_class_order}")
137
- self.cls_remap = np.array(
138
- [self.class_names.index(n) for n in model_class_order],
139
- dtype=np.int32,
140
- )
141
-
142
- for inp in self.session.get_inputs():
143
- print("INPUT:", inp.name, inp.shape, inp.type)
144
- for out in self.session.get_outputs():
145
- print("OUTPUT:", out.name, out.shape, out.type)
146
-
147
- self.input_name = self.session.get_inputs()[0].name
148
- self.output_names = [output.name for output in self.session.get_outputs()]
149
- self.input_shape = self.session.get_inputs()[0].shape
150
-
151
- # weights.onnx is exported at 1280x1280 (Ultralytics imgsz metadata),
152
- # static (dynamic=False). The default is only the fallback for when the
153
- # ONNX input dims aren't fixed; the real value is read from the session.
154
- self.input_height = self._safe_dim(self.input_shape[2], default=1280)
155
- self.input_width = self._safe_dim(self.input_shape[3], default=1280)
156
-
157
- self.use_tta = False
158
- self.use_tile_tta = False
159
- # Soft-NMS (ported from carwash001): Gaussian score decay of overlapping
160
- # boxes instead of hard removal. OFF by default to preserve the current
161
- # deployed behaviour; flip on (and tune sigma) via tune_miner.py to see if
162
- # it scores better — useful where signs cluster (gantries, sign assemblies).
163
- self.use_soft_nms = False
164
- self.soft_nms_sigma = 0.5
165
- self.soft_nms_score_thresh = 0.01
166
-
167
- print(f"✅ ONNX model loaded from: {model_path}")
168
- print(f"✅ ONNX providers: {self.session.get_providers()}")
169
- print(f"✅ ONNX input: name={self.input_name}, shape={self.input_shape}")
170
- print(f" ONNX input size: {self.input_width}x{self.input_height}, "
171
- f"use_tta={self.use_tta}, use_tile_tta={self.use_tile_tta}")
172
- print("per-class conf: " + ", ".join(
173
- f"{n}={t:.3f}" for n, t in zip(
174
- self.class_names, self._conf_thres_array.tolist()
175
- )
176
- ))
177
-
178
- self._warmup()
179
-
180
- def _warmup(self, iters: int = 3) -> None:
181
- try:
182
- dummy = np.zeros((720, 1280, 3), dtype=np.uint8)
183
- for _ in range(max(1, iters)):
184
- self.predict_batch(batch_images=[dummy], offset=0, n_keypoints=0)
185
- print(f"✅ warmup: {iters} dummy predict_batch call(s) done")
186
- except Exception as e:
187
- print(f"⚠️ warmup skipped: {e}")
188
-
189
- def _read_model_class_order(self) -> "list[str] | None":
190
- """Read the model's class order from Ultralytics ONNX metadata.
191
- Returns the class names ordered by model-emit index, or None when the
192
- metadata is missing/unparsable or doesn't match `class_names` as a set
193
- (in which case the static _model_class_order fallback is used)."""
194
- try:
195
- import ast
196
-
197
- meta = self.session.get_modelmeta().custom_metadata_map
198
- names = ast.literal_eval(meta["names"]) # e.g. {0: 'road_sign'}
199
- if isinstance(names, dict):
200
- order = [str(names[i]) for i in sorted(names)]
201
- else:
202
- order = [str(n) for n in names]
203
- except Exception as e:
204
- print(f"cls order: could not read ONNX names metadata ({e})")
205
- return None
206
- if sorted(order) != sorted(self.class_names):
207
- print(
208
- f"cls order: ONNX names {order} do not match expected classes "
209
- f"{self.class_names}; ignoring metadata"
210
- )
211
- return None
212
- return order
213
-
214
- def __repr__(self) -> str:
215
- return (
216
- f"ONNXRuntime(session={type(self.session).__name__}, "
217
- f"providers={self.session.get_providers()})"
218
- )
219
-
220
- @staticmethod
221
- def _safe_dim(value, default: int) -> int:
222
- return value if isinstance(value, int) and value > 0 else default
223
-
224
- def _letterbox(
225
- self,
226
- image: ndarray,
227
- new_shape: tuple[int, int],
228
- color=(114, 114, 114),
229
- ) -> tuple[ndarray, float, tuple[float, float]]:
230
- h, w = image.shape[:2]
231
- new_w, new_h = new_shape
232
-
233
- ratio = min(new_w / w, new_h / h)
234
- resized_w = int(round(w * ratio))
235
- resized_h = int(round(h * ratio))
236
-
237
- if (resized_w, resized_h) != (w, h):
238
- interp = cv2.INTER_CUBIC if ratio > 1.0 else cv2.INTER_LINEAR
239
- image = cv2.resize(image, (resized_w, resized_h), interpolation=interp)
240
-
241
- dw = (new_w - resized_w) / 2.0
242
- dh = (new_h - resized_h) / 2.0
243
-
244
- left = int(round(dw - 0.1))
245
- right = int(round(dw + 0.1))
246
- top = int(round(dh - 0.1))
247
- bottom = int(round(dh + 0.1))
248
-
249
- padded = cv2.copyMakeBorder(
250
- image, top, bottom, left, right,
251
- borderType=cv2.BORDER_CONSTANT, value=color,
252
- )
253
- return padded, ratio, (dw, dh)
254
-
255
- def _preprocess(
256
- self, image: ndarray
257
- ) -> tuple[np.ndarray, float, tuple[float, float], tuple[int, int]]:
258
- orig_h, orig_w = image.shape[:2]
259
- img, ratio, pad = self._letterbox(
260
- image, (self.input_width, self.input_height)
261
- )
262
- # Fused scale(1/255) + BGR->RGB swap + HWC->NCHW + contiguous float32 in
263
- # one optimized OpenCV call (bit-identical to the cvtColor + astype/255 +
264
- # transpose chain, but ~half the preprocess time).
265
- blob = cv2.dnn.blobFromImage(img, scalefactor=1.0 / 255.0, swapRB=True)
266
- return blob, ratio, pad, (orig_w, orig_h)
267
-
268
- @staticmethod
269
- def _clip_boxes(boxes: np.ndarray, image_size: tuple[int, int]) -> np.ndarray:
270
- w, h = image_size
271
- boxes[:, 0] = np.clip(boxes[:, 0], 0, w - 1)
272
- boxes[:, 1] = np.clip(boxes[:, 1], 0, h - 1)
273
- boxes[:, 2] = np.clip(boxes[:, 2], 0, w - 1)
274
- boxes[:, 3] = np.clip(boxes[:, 3], 0, h - 1)
275
- return boxes
276
-
277
- @staticmethod
278
- def _xywh_to_xyxy(boxes: np.ndarray) -> np.ndarray:
279
- out = np.empty_like(boxes)
280
- out[:, 0] = boxes[:, 0] - boxes[:, 2] / 2.0
281
- out[:, 1] = boxes[:, 1] - boxes[:, 3] / 2.0
282
- out[:, 2] = boxes[:, 0] + boxes[:, 2] / 2.0
283
- out[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.0
284
- return out
285
-
286
- @staticmethod
287
- def _hard_nms(
288
- boxes: np.ndarray, scores: np.ndarray, iou_thresh: float
289
- ) -> np.ndarray:
290
- n = len(boxes)
291
- if n == 0:
292
- return np.array([], dtype=np.intp)
293
- order = np.argsort(-scores)
294
- keep: list[int] = []
295
- while len(order) > 0:
296
- i = int(order[0])
297
- keep.append(i)
298
- if len(order) == 1:
299
- break
300
- rest = order[1:]
301
- xx1 = np.maximum(boxes[i, 0], boxes[rest, 0])
302
- yy1 = np.maximum(boxes[i, 1], boxes[rest, 1])
303
- xx2 = np.minimum(boxes[i, 2], boxes[rest, 2])
304
- yy2 = np.minimum(boxes[i, 3], boxes[rest, 3])
305
- inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
306
- a_i = (max(0.0, boxes[i, 2] - boxes[i, 0]) *
307
- max(0.0, boxes[i, 3] - boxes[i, 1]))
308
- a_r = (np.maximum(0.0, boxes[rest, 2] - boxes[rest, 0]) *
309
- np.maximum(0.0, boxes[rest, 3] - boxes[rest, 1]))
310
- iou = inter / (a_i + a_r - inter + 1e-7)
311
- order = rest[iou <= iou_thresh]
312
- return np.array(keep, dtype=np.intp)
313
-
314
- def _per_class_hard_nms(
315
- self,
316
- boxes: np.ndarray,
317
- scores: np.ndarray,
318
- cls_ids: np.ndarray,
319
- iou_thresh: float,
320
- ) -> np.ndarray:
321
- if len(boxes) == 0:
322
- return np.array([], dtype=np.intp)
323
- all_keep: list[int] = []
324
- for c in np.unique(cls_ids):
325
- mask = cls_ids == c
326
- indices = np.where(mask)[0]
327
- keep = self._hard_nms(boxes[mask], scores[mask], iou_thresh)
328
- all_keep.extend(indices[keep].tolist())
329
- all_keep.sort()
330
- return np.array(all_keep, dtype=np.intp)
331
-
332
- def _soft_nms(
333
- self,
334
- boxes: np.ndarray,
335
- scores: np.ndarray,
336
- sigma: float = 0.5,
337
- score_thresh: float = 0.01,
338
- ) -> tuple[np.ndarray, np.ndarray]:
339
- """Soft-NMS: Gaussian decay of overlapping scores instead of hard removal.
340
- Returns (kept_original_indices, updated_scores). (Ported from carwash001.)"""
341
- N = len(boxes)
342
- if N == 0:
343
- return np.array([], dtype=np.intp), np.array([], dtype=np.float32)
344
- boxes = boxes.astype(np.float32, copy=True)
345
- scores = scores.astype(np.float32, copy=True)
346
- order = np.arange(N)
347
- for i in range(N):
348
- max_pos = i + int(np.argmax(scores[i:]))
349
- boxes[[i, max_pos]] = boxes[[max_pos, i]]
350
- scores[[i, max_pos]] = scores[[max_pos, i]]
351
- order[[i, max_pos]] = order[[max_pos, i]]
352
- if i + 1 >= N:
353
- break
354
- xx1 = np.maximum(boxes[i, 0], boxes[i + 1:, 0])
355
- yy1 = np.maximum(boxes[i, 1], boxes[i + 1:, 1])
356
- xx2 = np.minimum(boxes[i, 2], boxes[i + 1:, 2])
357
- yy2 = np.minimum(boxes[i, 3], boxes[i + 1:, 3])
358
- inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
359
- area_i = max(0.0, float(
360
- (boxes[i, 2] - boxes[i, 0]) * (boxes[i, 3] - boxes[i, 1])))
361
- areas_j = (np.maximum(0.0, boxes[i + 1:, 2] - boxes[i + 1:, 0])
362
- * np.maximum(0.0, boxes[i + 1:, 3] - boxes[i + 1:, 1]))
363
- iou = inter / (area_i + areas_j - inter + 1e-7)
364
- scores[i + 1:] *= np.exp(-(iou ** 2) / sigma)
365
- mask = scores > score_thresh
366
- return order[mask], scores[mask]
367
-
368
- def _per_class_soft_nms(
369
- self,
370
- boxes: np.ndarray,
371
- scores: np.ndarray,
372
- cls_ids: np.ndarray,
373
- sigma: float = 0.5,
374
- score_thresh: float = 0.01,
375
- ) -> tuple[np.ndarray, np.ndarray]:
376
- """Soft-NMS applied independently per class. Returns (kept_idx, updated_scores)."""
377
- if len(boxes) == 0:
378
- return np.array([], dtype=np.intp), np.array([], dtype=np.float32)
379
- all_keep: list[int] = []
380
- all_scores: list[float] = []
381
- for c in np.unique(cls_ids):
382
- indices = np.where(cls_ids == c)[0]
383
- keep, updated = self._soft_nms(boxes[indices], scores[indices],
384
- sigma, score_thresh)
385
- for k, s in zip(keep, updated):
386
- all_keep.append(int(indices[k])); all_scores.append(float(s))
387
- if not all_keep:
388
- return np.array([], dtype=np.intp), np.array([], dtype=np.float32)
389
- return np.array(all_keep, dtype=np.intp), np.array(all_scores, dtype=np.float32)
390
-
391
- def _cross_class_dedup_op(
392
- self,
393
- boxes: np.ndarray,
394
- scores: np.ndarray,
395
- cls_ids: np.ndarray,
396
- iou_thresh: float,
397
- ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
398
- """Remove near-duplicate boxes across classes.
399
- Order candidates by (score - per_class_threshold) margin, then by area;
400
- keep the highest, suppress every other box with IoU > iou_thresh.
401
- With a single road_sign class this is effectively a no-op, but the
402
- method is kept so the pipeline stays compatible with the multi-class
403
- miner template.
404
- """
405
- n = len(boxes)
406
- if n <= 1:
407
- return boxes, scores, cls_ids
408
- boxes = np.asarray(boxes, dtype=np.float32)
409
- scores = np.asarray(scores, dtype=np.float32)
410
- cls_ids = np.asarray(cls_ids, dtype=np.int32)
411
- areas = (np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) *
412
- np.maximum(0.0, boxes[:, 3] - boxes[:, 1]))
413
- margins = scores - self._conf_thres_array[cls_ids]
414
- order = np.lexsort((-areas, -margins))
415
- suppressed = np.zeros(n, dtype=bool)
416
- keep: list[int] = []
417
- for i in order:
418
- if suppressed[i]:
419
- continue
420
- keep.append(int(i))
421
- bi = boxes[i]
422
- xx1 = np.maximum(bi[0], boxes[:, 0])
423
- yy1 = np.maximum(bi[1], boxes[:, 1])
424
- xx2 = np.minimum(bi[2], boxes[:, 2])
425
- yy2 = np.minimum(bi[3], boxes[:, 3])
426
- inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
427
- a_i = max(1e-7, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
428
- iou = inter / (a_i + areas - inter + 1e-7)
429
- dup = iou > iou_thresh
430
- dup[i] = False
431
- suppressed |= dup
432
- keep_idx = np.array(keep, dtype=np.intp)
433
- return boxes[keep_idx], scores[keep_idx], cls_ids[keep_idx]
434
-
435
- @staticmethod
436
- def _max_score_per_cluster(
437
- post_boxes: np.ndarray,
438
- post_cls: np.ndarray,
439
- full_boxes: np.ndarray,
440
- full_scores: np.ndarray,
441
- full_cls: np.ndarray,
442
- iou_thresh: float,
443
- ) -> np.ndarray:
444
- """For each kept (post-NMS) box, return the max score over the FULL
445
- candidate set among same-class boxes with IoU >= iou_thresh.
446
- Used after horizontal-flip TTA: a high-confidence flipped detection
447
- can raise the score of the corresponding original detection.
448
- """
449
- n = len(post_boxes)
450
- if n == 0:
451
- return np.empty(0, dtype=np.float32)
452
- full_areas = (np.maximum(0.0, full_boxes[:, 2] - full_boxes[:, 0]) *
453
- np.maximum(0.0, full_boxes[:, 3] - full_boxes[:, 1]))
454
- out = np.empty(n, dtype=np.float32)
455
- for i in range(n):
456
- bi = post_boxes[i]
457
- xx1 = np.maximum(bi[0], full_boxes[:, 0])
458
- yy1 = np.maximum(bi[1], full_boxes[:, 1])
459
- xx2 = np.minimum(bi[2], full_boxes[:, 2])
460
- yy2 = np.minimum(bi[3], full_boxes[:, 3])
461
- inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
462
- a_i = max(0.0, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
463
- iou = inter / (a_i + full_areas - inter + 1e-7)
464
- cluster = (iou >= iou_thresh) & (full_cls == post_cls[i])
465
- out[i] = float(np.max(full_scores[cluster])) if np.any(cluster) else 0.0
466
- return out
467
-
468
- def _conf_filter_mask(
469
- self, scores: np.ndarray, cls_ids: np.ndarray
470
- ) -> np.ndarray:
471
- """Boolean keep-mask: score >= per-class threshold, with a per-class
472
- rescue -- if a class has zero boxes passing, admit its top-1 candidate
473
- when its score >= (per-class threshold - per-class bonus)."""
474
- if len(scores) == 0:
475
- return np.zeros(0, dtype=bool)
476
- thr = self._conf_thres_array[cls_ids]
477
- keep = scores >= thr
478
- for c in np.unique(cls_ids):
479
- b = float(self._bonus_array[c])
480
- if b <= 0.0:
481
- continue
482
- cm = cls_ids == c
483
- if keep[cm].any():
484
- continue
485
- idx = np.where(cm)[0]
486
- top = int(idx[int(np.argmax(scores[idx]))])
487
- if scores[top] >= self._conf_thres_array[c] - b:
488
- keep[top] = True
489
- return keep
490
-
491
- def _filter_sane_boxes(
492
- self,
493
- boxes: np.ndarray,
494
- scores: np.ndarray,
495
- cls_ids: np.ndarray,
496
- orig_size: tuple[int, int],
497
- ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
498
- """Drop tiny / degenerate / image-spanning / extreme-AR boxes (FP)."""
499
- if len(boxes) == 0:
500
- return boxes, scores, cls_ids
501
- orig_w, orig_h = orig_size
502
- image_area = float(orig_w * orig_h)
503
- keep = []
504
- for i, box in enumerate(boxes):
505
- x1, y1, x2, y2 = box.tolist()
506
- bw = x2 - x1
507
- bh = y2 - y1
508
- if bw <= 0 or bh <= 0:
509
- continue
510
- if bw < self.min_side or bh < self.min_side:
511
- continue
512
- area = bw * bh
513
- if area < self.min_box_area:
514
- continue
515
- if area > 0.95 * image_area:
516
- continue
517
- ar = max(bw / max(bh, 1e-6), bh / max(bw, 1e-6))
518
- if ar > self.max_aspect_ratio:
519
- continue
520
- keep.append(i)
521
- if not keep:
522
- return (
523
- np.empty((0, 4), dtype=np.float32),
524
- np.empty((0,), dtype=np.float32),
525
- np.empty((0,), dtype=np.int32),
526
- )
527
- k = np.array(keep, dtype=np.intp)
528
- return boxes[k], scores[k], cls_ids[k]
529
-
530
- def _per_view_pipeline(
531
- self,
532
- boxes: np.ndarray,
533
- scores: np.ndarray,
534
- cls_ids: np.ndarray,
535
- ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
536
- """Per-view post-processing pipeline: per-class NMS -> cap -> cross-class dedup."""
537
- if len(boxes) > 1:
538
- if self.use_soft_nms:
539
- keep, new_scores = self._per_class_soft_nms(
540
- boxes, scores, cls_ids,
541
- self.soft_nms_sigma, self.soft_nms_score_thresh)
542
- boxes, scores, cls_ids = boxes[keep], new_scores, cls_ids[keep]
543
- else:
544
- keep = self._per_class_hard_nms(boxes, scores, cls_ids, self.iou_thres)
545
- boxes, scores, cls_ids = boxes[keep], scores[keep], cls_ids[keep]
546
- if len(scores) > self.max_det:
547
- top = np.argsort(-scores)[: self.max_det]
548
- boxes, scores, cls_ids = boxes[top], scores[top], cls_ids[top]
549
- if len(boxes) > 1:
550
- boxes, scores, cls_ids = self._cross_class_dedup_op(
551
- boxes, scores, cls_ids, self.cross_iou_thresh
552
- )
553
- return boxes, scores, cls_ids
554
-
555
- @staticmethod
556
- def _shrink_wh(coords: np.ndarray, sw: float, sh: float) -> np.ndarray:
557
- """Shrink each xyxy box about its center: new_w = w/sw, new_h = h/sh."""
558
- if len(coords) == 0:
559
- return coords
560
- coords = np.asarray(coords, dtype=np.float32).copy()
561
- cx = (coords[:, 0] + coords[:, 2]) * 0.5
562
- cy = (coords[:, 1] + coords[:, 3]) * 0.5
563
- hw = (coords[:, 2] - coords[:, 0]) * (0.5 / sw)
564
- hh = (coords[:, 3] - coords[:, 1]) * (0.5 / sh)
565
- coords[:, 0] = cx - hw
566
- coords[:, 1] = cy - hh
567
- coords[:, 2] = cx + hw
568
- coords[:, 3] = cy + hh
569
- return coords
570
-
571
- def _build_results(
572
- self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray
573
- ) -> list[BoundingBox]:
574
- boxes = self._shrink_wh(boxes, self.box_shrink_w, self.box_shrink_h)
575
- results: list[BoundingBox] = []
576
- for box, conf, cls_id in zip(boxes, scores, cls_ids):
577
- x1, y1, x2, y2 = box.tolist()
578
- if x2 <= x1 or y2 <= y1:
579
- continue
580
- results.append(
581
- BoundingBox(
582
- x1=int(round(x1)),
583
- y1=int(round(y1)),
584
- x2=int(round(x2)),
585
- y2=int(round(y2)),
586
- cls_id=int(cls_id),
587
- conf=float(conf),
588
- )
589
- )
590
- return results
591
-
592
- def _decode_final_dets(
593
- self,
594
- preds: np.ndarray,
595
- ratio: float,
596
- pad: tuple[float, float],
597
- orig_size: tuple[int, int],
598
- ) -> list[BoundingBox]:
599
- """Final-detection output path: rows shaped [x1, y1, x2, y2, conf, cls_id]."""
600
- if preds.ndim == 3 and preds.shape[0] == 1:
601
- preds = preds[0]
602
- if preds.ndim != 2 or preds.shape[1] < 6:
603
- raise ValueError(f"Unexpected ONNX final-det output shape: {preds.shape}")
604
-
605
- boxes = preds[:, :4].astype(np.float32)
606
- scores = preds[:, 4].astype(np.float32)
607
- cls_ids = preds[:, 5].astype(np.int32)
608
- cls_ids = self.cls_remap[cls_ids]
609
-
610
- keep = self._conf_filter_mask(scores, cls_ids)
611
- boxes = boxes[keep]
612
- scores = scores[keep]
613
- cls_ids = cls_ids[keep]
614
- if len(boxes) == 0:
615
- return []
616
-
617
- pad_w, pad_h = pad
618
- boxes[:, [0, 2]] -= pad_w
619
- boxes[:, [1, 3]] -= pad_h
620
- boxes /= ratio
621
- boxes = self._clip_boxes(boxes, orig_size)
622
-
623
- boxes, scores, cls_ids = self._filter_sane_boxes(
624
- boxes, scores, cls_ids, orig_size
625
- )
626
- if len(boxes) == 0:
627
- return []
628
-
629
- boxes, scores, cls_ids = self._per_view_pipeline(boxes, scores, cls_ids)
630
- return self._build_results(boxes, scores, cls_ids)
631
-
632
- def _decode_raw_yolo(
633
- self,
634
- preds: np.ndarray,
635
- ratio: float,
636
- pad: tuple[float, float],
637
- orig_size: tuple[int, int],
638
- ) -> list[BoundingBox]:
639
- """Fallback raw-YOLO output path: per-anchor class logits."""
640
- if preds.ndim != 3 or preds.shape[0] != 1:
641
- raise ValueError(f"Unexpected raw ONNX output shape: {preds.shape}")
642
- preds = preds[0]
643
- if preds.shape[0] <= 16 and preds.shape[1] > preds.shape[0]:
644
- preds = preds.T
645
- if preds.ndim != 2 or preds.shape[1] < 5:
646
- raise ValueError(f"Unexpected raw output shape: {preds.shape}")
647
-
648
- boxes_xywh = preds[:, :4].astype(np.float32)
649
- cls_part = preds[:, 4:].astype(np.float32)
650
- if cls_part.shape[1] == 1:
651
- scores = cls_part[:, 0]
652
- cls_ids = np.zeros(len(scores), dtype=np.int32)
653
- else:
654
- cls_ids = np.argmax(cls_part, axis=1).astype(np.int32)
655
- scores = cls_part[np.arange(len(cls_part)), cls_ids]
656
- cls_ids = self.cls_remap[cls_ids]
657
-
658
- keep = self._conf_filter_mask(scores, cls_ids)
659
- boxes_xywh = boxes_xywh[keep]
660
- scores = scores[keep]
661
- cls_ids = cls_ids[keep]
662
- if len(boxes_xywh) == 0:
663
- return []
664
- boxes = self._xywh_to_xyxy(boxes_xywh)
665
-
666
- pad_w, pad_h = pad
667
- boxes[:, [0, 2]] -= pad_w
668
- boxes[:, [1, 3]] -= pad_h
669
- boxes /= ratio
670
- boxes = self._clip_boxes(boxes, orig_size)
671
-
672
- boxes, scores, cls_ids = self._filter_sane_boxes(
673
- boxes, scores, cls_ids, orig_size
674
- )
675
- if len(boxes) == 0:
676
- return []
677
-
678
- boxes, scores, cls_ids = self._per_view_pipeline(boxes, scores, cls_ids)
679
- return self._build_results(boxes, scores, cls_ids)
680
-
681
- def _postprocess(
682
- self,
683
- output: np.ndarray,
684
- ratio: float,
685
- pad: tuple[float, float],
686
- orig_size: tuple[int, int],
687
- ) -> list[BoundingBox]:
688
- if output.ndim == 2 and output.shape[1] >= 6:
689
- return self._decode_final_dets(output, ratio, pad, orig_size)
690
- if output.ndim == 3 and output.shape[0] == 1 and output.shape[2] == 6:
691
- return self._decode_final_dets(output, ratio, pad, orig_size)
692
- return self._decode_raw_yolo(output, ratio, pad, orig_size)
693
-
694
- def _predict_single(self, image: np.ndarray) -> list[BoundingBox]:
695
- if image is None:
696
- raise ValueError("Input image is None")
697
- if not isinstance(image, np.ndarray):
698
- raise TypeError(f"Input is not numpy array: {type(image)}")
699
- if image.ndim != 3:
700
- raise ValueError(f"Expected HWC image, got shape={image.shape}")
701
- if image.shape[0] <= 0 or image.shape[1] <= 0:
702
- raise ValueError(f"Invalid image shape={image.shape}")
703
- if image.shape[2] != 3:
704
- raise ValueError(f"Expected 3 channels, got shape={image.shape}")
705
- if image.dtype != np.uint8:
706
- image = image.astype(np.uint8)
707
-
708
- input_tensor, ratio, pad, orig_size = self._preprocess(image)
709
- expected = (1, 3, self.input_height, self.input_width)
710
- if input_tensor.shape != expected:
711
- raise ValueError(
712
- f"Bad input tensor shape={input_tensor.shape}, expected={expected}"
713
- )
714
-
715
- outputs = self.session.run(self.output_names, {self.input_name: input_tensor})
716
- return self._postprocess(outputs[0], ratio, pad, orig_size)
717
-
718
- def _predict_tta(self, image: np.ndarray) -> list[BoundingBox]:
719
- """Horizontal-flip TTA.
720
- Strategy:
721
- 1. Predict on original and on flipped image.
722
- 2. Map flipped boxes back to original coordinates.
723
- 3. Per-class hard NMS on the union.
724
- 4. For each kept box, compute the max same-class score across the
725
- FULL union (not just the post-NMS subset) -- this lets a high-
726
- confidence flipped detection raise a borderline original one.
727
- 5. Cross-class dedup to suppress same-physical-object multi-class.
728
- """
729
- boxes_orig = self._predict_single(image)
730
- flipped = cv2.flip(image, 1)
731
- boxes_flip = self._predict_single(flipped)
732
- w = image.shape[1]
733
- boxes_flip = [
734
- BoundingBox(
735
- x1=w - b.x2, y1=b.y1, x2=w - b.x1, y2=b.y2,
736
- cls_id=b.cls_id, conf=b.conf,
737
- )
738
- for b in boxes_flip
739
- ]
740
- all_boxes = boxes_orig + boxes_flip
741
- if not all_boxes:
742
- return []
743
-
744
- coords = np.array(
745
- [[b.x1, b.y1, b.x2, b.y2] for b in all_boxes], dtype=np.float32
746
- )
747
- scores = np.array([b.conf for b in all_boxes], dtype=np.float32)
748
- cls_ids = np.array([b.cls_id for b in all_boxes], dtype=np.int32)
749
-
750
- hard_keep = self._per_class_hard_nms(coords, scores, cls_ids, self.iou_thres)
751
- if len(hard_keep) == 0:
752
- return []
753
- if len(hard_keep) > self.max_det:
754
- top = np.argsort(-scores[hard_keep])[: self.max_det]
755
- hard_keep = hard_keep[top]
756
-
757
- boosted = self._max_score_per_cluster(
758
- coords[hard_keep], cls_ids[hard_keep],
759
- coords, scores, cls_ids, self.iou_thres,
760
- )
761
-
762
- kept_coords = coords[hard_keep]
763
- kept_cls = cls_ids[hard_keep]
764
- if len(kept_coords) > 1:
765
- kept_coords, boosted, kept_cls = self._cross_class_dedup_op(
766
- kept_coords, boosted, kept_cls, self.cross_iou_thresh
767
- )
768
-
769
- kept_coords = self._shrink_wh(
770
- kept_coords, self.box_shrink_w, self.box_shrink_h
771
- )
772
- return [
773
- BoundingBox(
774
- x1=int(round(float(kept_coords[j, 0]))),
775
- y1=int(round(float(kept_coords[j, 1]))),
776
- x2=int(round(float(kept_coords[j, 2]))),
777
- y2=int(round(float(kept_coords[j, 3]))),
778
- cls_id=int(kept_cls[j]),
779
- conf=float(boosted[j]),
780
- )
781
- for j in range(len(kept_coords))
782
- ]
783
-
784
- def _predict_tiles(self, image: np.ndarray) -> list[BoundingBox]:
785
- """Tile-based TTA for high-resolution images.
786
- Splits the source image into two overlapping horizontal tiles, runs
787
- single-pass inference on each at native scale, and translates boxes
788
- back to the global frame. Useful when source width >> model input
789
- width because letterboxing otherwise discards effective resolution
790
- that small / distant signs depend on.
791
- Returns an empty list if the image isn't wide enough to benefit; the
792
- caller falls back to the regular pipeline in that case.
793
- """
794
- h, w = image.shape[:2]
795
- if w < int(self.input_width * self.tile_trigger_ratio):
796
- return []
797
-
798
- overlap = int(w * self.tile_overlap_ratio)
799
- mid = w // 2
800
- x_left_end = min(w, mid + overlap // 2)
801
- x_right_start = max(0, mid - overlap // 2)
802
-
803
- left = image[:, :x_left_end]
804
- right = image[:, x_right_start:]
805
-
806
- boxes_left = self._predict_single(left)
807
- boxes_right = self._predict_single(right)
808
-
809
- shifted_right = [
810
- BoundingBox(
811
- x1=b.x1 + x_right_start,
812
- y1=b.y1,
813
- x2=b.x2 + x_right_start,
814
- y2=b.y2,
815
- cls_id=b.cls_id,
816
- conf=b.conf,
817
- )
818
- for b in boxes_right
819
- ]
820
- return boxes_left + shifted_right
821
-
822
- def _merge_views(
823
- self,
824
- view_boxes: list[list[BoundingBox]],
825
- image_size: tuple[int, int],
826
- ) -> list[BoundingBox]:
827
- """Merge boxes from multiple views (single / hflip / tiles).
828
- Same logic as `_predict_tta`'s tail: per-class hard NMS to dedupe,
829
- then for each kept box take the max same-class score across the full
830
- candidate union a high-confidence detection in any view boosts
831
- borderline matches in others.
832
- """
833
- all_boxes: list[BoundingBox] = []
834
- for vb in view_boxes:
835
- all_boxes.extend(vb)
836
- if not all_boxes:
837
- return []
838
-
839
- coords = np.array(
840
- [[b.x1, b.y1, b.x2, b.y2] for b in all_boxes], dtype=np.float32
841
- )
842
- scores = np.array([b.conf for b in all_boxes], dtype=np.float32)
843
- cls_ids = np.array([b.cls_id for b in all_boxes], dtype=np.int32)
844
-
845
- coords = self._clip_boxes(coords, image_size)
846
-
847
- hard_keep = self._per_class_hard_nms(coords, scores, cls_ids, self.iou_thres)
848
- if len(hard_keep) == 0:
849
- return []
850
- if len(hard_keep) > self.max_det:
851
- top = np.argsort(-scores[hard_keep])[: self.max_det]
852
- hard_keep = hard_keep[top]
853
-
854
- boosted = self._max_score_per_cluster(
855
- coords[hard_keep], cls_ids[hard_keep],
856
- coords, scores, cls_ids, self.iou_thres,
857
- )
858
-
859
- kept_coords = coords[hard_keep]
860
- kept_cls = cls_ids[hard_keep]
861
- if len(kept_coords) > 1:
862
- kept_coords, boosted, kept_cls = self._cross_class_dedup_op(
863
- kept_coords, boosted, kept_cls, self.cross_iou_thresh
864
- )
865
-
866
- kept_coords = self._shrink_wh(
867
- kept_coords, self.box_shrink_w, self.box_shrink_h
868
- )
869
- return [
870
- BoundingBox(
871
- x1=int(round(float(kept_coords[j, 0]))),
872
- y1=int(round(float(kept_coords[j, 1]))),
873
- x2=int(round(float(kept_coords[j, 2]))),
874
- y2=int(round(float(kept_coords[j, 3]))),
875
- cls_id=int(kept_cls[j]),
876
- conf=float(boosted[j]),
877
- )
878
- for j in range(len(kept_coords))
879
- ]
880
-
881
- def _predict_full(self, image: np.ndarray) -> list[BoundingBox]:
882
- """Top-level per-frame prediction with all enabled augmentations.
883
- - `use_tta=True`: original + horizontal flip
884
- - `use_tile_tta=True` AND image wide enough: two overlapping tiles
885
- All views are merged via per-class NMS + cluster-max score boost.
886
- """
887
- if not self.use_tta and not self.use_tile_tta:
888
- return self._predict_single(image)
889
-
890
- views: list[list[BoundingBox]] = []
891
- if self.use_tta:
892
- views.append(self._predict_single(image))
893
- flipped = cv2.flip(image, 1)
894
- w = image.shape[1]
895
- flipped_dets = self._predict_single(flipped)
896
- views.append([
897
- BoundingBox(
898
- x1=w - b.x2, y1=b.y1, x2=w - b.x1, y2=b.y2,
899
- cls_id=b.cls_id, conf=b.conf,
900
- )
901
- for b in flipped_dets
902
- ])
903
- else:
904
- views.append(self._predict_single(image))
905
-
906
- if self.use_tile_tta:
907
- tile_boxes = self._predict_tiles(image)
908
- if tile_boxes:
909
- views.append(tile_boxes)
910
-
911
- h, w = image.shape[:2]
912
- return self._merge_views(views, (w, h))
913
-
914
- def predict_batch(
915
- self,
916
- batch_images: list[ndarray],
917
- offset: int,
918
- n_keypoints: int,
919
- ) -> list[TVFrameResult]:
920
- results: list[TVFrameResult] = []
921
- for frame_number_in_batch, image in enumerate(batch_images):
922
- try:
923
- boxes = self._predict_full(image)
924
- except Exception as e:
925
- print(
926
- f"⚠️ Inference failed for frame "
927
- f"{offset + frame_number_in_batch}: {e}"
928
- )
929
- boxes = []
930
- results.append(
931
- TVFrameResult(
932
- frame_id=offset + frame_number_in_batch,
933
- boxes=boxes,
934
- keypoints=[(0, 0) for _ in range(max(0, int(n_keypoints)))],
935
- )
936
- )
937
- return results
 
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
+
11
+ class BoundingBox(BaseModel):
12
+ x1: int
13
+ y1: int
14
+ x2: int
15
+ y2: int
16
+ cls_id: int
17
+ conf: float
18
+
19
+
20
+ class TVFrameResult(BaseModel):
21
+ frame_id: int
22
+ boxes: list[BoundingBox]
23
+ keypoints: list[tuple[int, int]]
24
+
25
+
26
+ class Miner:
27
+ def __init__(self,
28
+ path_hf_repo: Path
29
+ ) -> None:
30
+ model_path = self._resolve_model_path(path_hf_repo)
31
+ # car-wash element classes cls_id order MUST match element `objects`
32
+ # (0=broom, 1=drainage gate, 2=nozzle, 3=track). This is the canonical
33
+ # order every downstream consumer (validator, BoundingBox.cls_id) sees.
34
+ self.class_names = ["broom", "drainage gate", "nozzle", "track"]
35
+ # FALLBACK model-emit order: the authoritative order is read from the
36
+ # ONNX `names` metadata after the session is created (embedded by
37
+ # Ultralytics at export, ships inside weights.onnx), so a retrained
38
+ # model with a different class order is remapped correctly without
39
+ # code changes. This list is used only when metadata is missing.
40
+ self._model_class_order = ["broom", "drainage gate", "nozzle", "track"]
41
+ print("ORT version:", ort.__version__)
42
+
43
+ try:
44
+ ort.preload_dlls()
45
+ print("✅ onnxruntime.preload_dlls() success")
46
+ except Exception as e:
47
+ print(f"⚠️ preload_dlls failed: {e}")
48
+
49
+ print("ORT available providers BEFORE session:", ort.get_available_providers())
50
+
51
+ sess_options = ort.SessionOptions()
52
+ sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
53
+ sess_options.intra_op_num_threads = 2
54
+ sess_options.inter_op_num_threads = 1
55
+ sess_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
56
+
57
+ try:
58
+ self.session = ort.InferenceSession(
59
+ str(model_path),
60
+ sess_options=sess_options,
61
+ providers=["CPUExecutionProvider"],
62
+ )
63
+ except Exception as e:
64
+ self.session = ort.InferenceSession(
65
+ str(model_path),
66
+ sess_options=sess_options,
67
+ providers=["CPUExecutionProvider"],
68
+ )
69
+
70
+ print("ORT session providers:", self.session.get_providers())
71
+
72
+ # Build cls_remap: for each model-emit index i,
73
+ # cls_remap[i] = self.class_names.index(model_class_order[i])
74
+ # The model-side order comes from the ONNX metadata when available,
75
+ # else falls back to the static _model_class_order.
76
+ model_class_order = self._read_model_class_order()
77
+ if model_class_order is None:
78
+ model_class_order = list(self._model_class_order)
79
+ print(f"cls order: no usable ONNX metadata, FALLBACK {model_class_order}")
80
+ else:
81
+ print(f"cls order: from ONNX metadata {model_class_order}")
82
+ self.cls_remap = np.array(
83
+ [self.class_names.index(n) for n in model_class_order], dtype=np.int32
84
+ )
85
+
86
+ for inp in self.session.get_inputs():
87
+ print("INPUT:", inp.name, inp.shape, inp.type)
88
+
89
+ for out in self.session.get_outputs():
90
+ print("OUTPUT:", out.name, out.shape, out.type)
91
+
92
+ self.input_name = self.session.get_inputs()[0].name
93
+ self.output_names = [output.name for output in self.session.get_outputs()]
94
+ self.input_shape = self.session.get_inputs()[0].shape
95
+
96
+ # Match the ONNX input dtype (this export is FP16 -> needs float16 input).
97
+ input_type = self.session.get_inputs()[0].type
98
+ self.np_dtype = np.float16 if "float16" in input_type else np.float32
99
+ print(f"✅ ONNX input dtype: {input_type} -> numpy {self.np_dtype}")
100
+
101
+ # ONNX is fixed-size 1408x1408 (v1 export); read actual shape to be safe.
102
+ self.input_height = self._safe_dim(self.input_shape[2], default=1280)
103
+ self.input_width = self._safe_dim(self.input_shape[3], default=1280)
104
+
105
+ # Tuned for validator scoring (pillars: 0.6*map50 + 0.4*false_positive).
106
+ # All values below are the measured optimum of a full TTA-off grid sweep
107
+ # on the validator-style val split (tune_miner.py, car-wash-49-val1024
108
+ # val, ALL 2476 crops, against the synthetic-crop model; composite
109
+ # 0.8716 -> 0.8731) -- re-run the sweep after any retrain.
110
+ self.iou_thres = 0.5 # Per-class NMS IoU; lower = stricter dedup
111
+ self.cross_iou_thresh = 0.9 # Cross-class dedup IoU (suppress same physical object firing multiple classes)
112
+ self.max_det = 200
113
+ # TTA = a 2nd (flipped) forward pass. Doubles latency; off for the
114
+ # CPU latency gate. Re-enable only if the latency budget allows.
115
+ self.use_tta = False
116
+
117
+ # conf thresholds: broom=0.38 drainage gate=0.45 nozzle=0.30 track=0.60
118
+ # Per-class confidence thresholds.
119
+ # Indexed by class_names order: [broom, drainage gate, nozzle, track].
120
+ # Values are the 2476-crop TTA-off sweep optimum on the synthetic-crop
121
+ # model: lowering drainage (0.30->0.22) and nozzle (0.45->0.38) recovers
122
+ # recall (map50 0.836->0.847) while FP actually drops (0.925->0.913).
123
+ self._conf_thres_array = np.array(
124
+ [0.35, 0.30, 0.45, 0.30], dtype=np.float32
125
+ )
126
+ # Per-class rescue bonus: when a class has ZERO boxes passing the
127
+ # threshold in a frame, its top-1 candidate is admitted when its score
128
+ # is at least (per-class threshold - per-class bonus).
129
+ # DISABLED (all zeros): the 2476-crop sweep (rescue_bonus=False won)
130
+ # confirmed rescue admits more false positives than true positives
131
+ # under the validator's FP pillar.
132
+ self._bonus_array = np.array(
133
+ [0.05, 0.05, 0.1, 0.05], dtype=np.float32
134
+ )
135
+
136
+ # Box sanity filter — kept loose: car-wash `nozzle` boxes are tiny
137
+ # (GT median ~290 px², smallest ~32 px²). Fire's 14x14/min_side 8
138
+ # would delete valid nozzles, so thresholds are dropped here.
139
+ self.min_box_area = 4 * 4 # 16 px²
140
+ self.min_side = 3
141
+ self.max_aspect_ratio = 12.0
142
+
143
+ print(f" ONNX model loaded from: {model_path}")
144
+ print(f"✅ ONNX providers: {self.session.get_providers()}")
145
+ print(f"✅ ONNX input: name={self.input_name}, shape={self.input_shape}")
146
+
147
+ self._warmup()
148
+
149
+ def _warmup(self, iters: int = 3) -> None:
150
+ try:
151
+ dummy = np.zeros((720, 1280, 3), dtype=np.uint8)
152
+ for _ in range(max(1, iters)):
153
+ self.predict_batch(batch_images=[dummy], offset=0, n_keypoints=0)
154
+ print(f"✅ warmup: {iters} dummy predict_batch call(s) done")
155
+ except Exception as e:
156
+ print(f"⚠️ warmup skipped: {e}")
157
+
158
+ def __repr__(self) -> str:
159
+ return (
160
+ f"ONNXRuntime(session={type(self.session).__name__}, "
161
+ f"providers={self.session.get_providers()})"
162
+ )
163
+
164
+ @staticmethod
165
+ def _safe_dim(value, default: int) -> int:
166
+ return value if isinstance(value, int) and value > 0 else default
167
+
168
+ @staticmethod
169
+ def _resolve_model_path(repo: Path) -> Path:
170
+ """Locate the ONNX model in the repo dir.
171
+
172
+ Prefers weights.onnx (FP16/FP32 export), then weights_int8.onnx (the
173
+ training script's INT8-quantized export -- works as-is: quantization
174
+ preserves the Ultralytics metadata and QDQ models take regular fp32
175
+ input), then any other .onnx file. INT8 is the fallback when the FP16
176
+ export exceeds the 30 MB deployment limit (e.g. yolo26m).
177
+ """
178
+ for name in ("weights.onnx", "weights_int8.onnx"):
179
+ p = repo / name
180
+ if p.exists():
181
+ if name != "weights.onnx":
182
+ print(f"model: weights.onnx not found, using {name}")
183
+ return p
184
+ candidates = sorted(repo.glob("*.onnx"))
185
+ if candidates:
186
+ print(f"model: using {candidates[0].name}")
187
+ return candidates[0]
188
+ return repo / "weights.onnx" # let session creation raise the error
189
+
190
+ def _read_model_class_order(self) -> list[str] | None:
191
+ """Read the model's class order from Ultralytics ONNX metadata.
192
+
193
+ Returns the class names ordered by model-emit index, or None when
194
+ metadata is missing/unparsable or doesn't match `class_names` as a
195
+ set (in which case the static _model_class_order fallback is used).
196
+ """
197
+ try:
198
+ import ast
199
+
200
+ meta = self.session.get_modelmeta().custom_metadata_map
201
+ names = ast.literal_eval(meta["names"]) # e.g. {0: 'broom', ...}
202
+ if isinstance(names, dict):
203
+ order = [str(names[i]) for i in sorted(names)]
204
+ else:
205
+ order = [str(n) for n in names]
206
+ except Exception as e:
207
+ print(f"cls order: could not read ONNX names metadata ({e})")
208
+ return None
209
+ if sorted(order) != sorted(self.class_names):
210
+ print(
211
+ f"cls order: ONNX names {order} do not match expected classes "
212
+ f"{self.class_names}; ignoring metadata"
213
+ )
214
+ return None
215
+ return order
216
+
217
+ def _letterbox(
218
+ self,
219
+ image: ndarray,
220
+ new_shape: tuple[int, int],
221
+ color=(114, 114, 114),
222
+ ) -> tuple[ndarray, float, tuple[float, float]]:
223
+ """
224
+ Resize with unchanged aspect ratio and pad to target shape.
225
+ Returns:
226
+ padded_image,
227
+ ratio,
228
+ (pad_w, pad_h) # half-padding
229
+ """
230
+ h, w = image.shape[:2]
231
+ new_w, new_h = new_shape
232
+
233
+ ratio = min(new_w / w, new_h / h)
234
+ resized_w = int(round(w * ratio))
235
+ resized_h = int(round(h * ratio))
236
+
237
+ if (resized_w, resized_h) != (w, h):
238
+ interp = cv2.INTER_CUBIC if ratio > 1.0 else cv2.INTER_LINEAR
239
+ image = cv2.resize(image, (resized_w, resized_h), interpolation=interp)
240
+
241
+ dw = new_w - resized_w
242
+ dh = new_h - resized_h
243
+ dw /= 2.0
244
+ dh /= 2.0
245
+
246
+ left = int(round(dw - 0.1))
247
+ right = int(round(dw + 0.1))
248
+ top = int(round(dh - 0.1))
249
+ bottom = int(round(dh + 0.1))
250
+
251
+ padded = cv2.copyMakeBorder(
252
+ image,
253
+ top,
254
+ bottom,
255
+ left,
256
+ right,
257
+ borderType=cv2.BORDER_CONSTANT,
258
+ value=color,
259
+ )
260
+ return padded, ratio, (dw, dh)
261
+
262
+ def _preprocess(
263
+ self, image: ndarray
264
+ ) -> tuple[np.ndarray, float, tuple[float, float], tuple[int, int]]:
265
+ """
266
+ Preprocess for fixed-size ONNX export:
267
+ - enhance image quality (CLAHE, denoise, sharpen)
268
+ - letterbox to model input size
269
+ - BGR -> RGB
270
+ - normalize to [0,1]
271
+ - HWC -> NCHW float32
272
+ """
273
+ orig_h, orig_w = image.shape[:2]
274
+
275
+ img, ratio, pad = self._letterbox(
276
+ image, (self.input_width, self.input_height)
277
+ )
278
+ img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
279
+ img = (img.astype(np.float32) / 255.0)
280
+ img = np.transpose(img, (2, 0, 1))[None, ...]
281
+ img = np.ascontiguousarray(img, dtype=self.np_dtype)
282
+
283
+ return img, ratio, pad, (orig_w, orig_h)
284
+
285
+ @staticmethod
286
+ def _clip_boxes(boxes: np.ndarray, image_size: tuple[int, int]) -> np.ndarray:
287
+ w, h = image_size
288
+ boxes[:, 0] = np.clip(boxes[:, 0], 0, w - 1)
289
+ boxes[:, 1] = np.clip(boxes[:, 1], 0, h - 1)
290
+ boxes[:, 2] = np.clip(boxes[:, 2], 0, w - 1)
291
+ boxes[:, 3] = np.clip(boxes[:, 3], 0, h - 1)
292
+ return boxes
293
+
294
+ @staticmethod
295
+ def _xywh_to_xyxy(boxes: np.ndarray) -> np.ndarray:
296
+ out = np.empty_like(boxes)
297
+ out[:, 0] = boxes[:, 0] - boxes[:, 2] / 2.0
298
+ out[:, 1] = boxes[:, 1] - boxes[:, 3] / 2.0
299
+ out[:, 2] = boxes[:, 0] + boxes[:, 2] / 2.0
300
+ out[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.0
301
+ return out
302
+
303
+ def _soft_nms(
304
+ self,
305
+ boxes: np.ndarray,
306
+ scores: np.ndarray,
307
+ sigma: float = 0.5,
308
+ score_thresh: float = 0.01,
309
+ ) -> tuple[np.ndarray, np.ndarray]:
310
+ """
311
+ Soft-NMS: Gaussian decay of overlapping scores instead of hard removal.
312
+ Returns (kept_original_indices, updated_scores).
313
+ """
314
+ N = len(boxes)
315
+ if N == 0:
316
+ return np.array([], dtype=np.intp), np.array([], dtype=np.float32)
317
+
318
+ boxes = boxes.astype(np.float32, copy=True)
319
+ scores = scores.astype(np.float32, copy=True)
320
+ order = np.arange(N)
321
+
322
+ for i in range(N):
323
+ max_pos = i + int(np.argmax(scores[i:]))
324
+ boxes[[i, max_pos]] = boxes[[max_pos, i]]
325
+ scores[[i, max_pos]] = scores[[max_pos, i]]
326
+ order[[i, max_pos]] = order[[max_pos, i]]
327
+
328
+ if i + 1 >= N:
329
+ break
330
+
331
+ xx1 = np.maximum(boxes[i, 0], boxes[i + 1:, 0])
332
+ yy1 = np.maximum(boxes[i, 1], boxes[i + 1:, 1])
333
+ xx2 = np.minimum(boxes[i, 2], boxes[i + 1:, 2])
334
+ yy2 = np.minimum(boxes[i, 3], boxes[i + 1:, 3])
335
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
336
+
337
+ area_i = max(0.0, float(
338
+ (boxes[i, 2] - boxes[i, 0]) * (boxes[i, 3] - boxes[i, 1])
339
+ ))
340
+ areas_j = (
341
+ np.maximum(0.0, boxes[i + 1:, 2] - boxes[i + 1:, 0])
342
+ * np.maximum(0.0, boxes[i + 1:, 3] - boxes[i + 1:, 1])
343
+ )
344
+ iou = inter / (area_i + areas_j - inter + 1e-7)
345
+ scores[i + 1:] *= np.exp(-(iou ** 2) / sigma)
346
+
347
+ mask = scores > score_thresh
348
+ return order[mask], scores[mask]
349
+
350
+ @staticmethod
351
+ def _hard_nms(
352
+ boxes: np.ndarray,
353
+ scores: np.ndarray,
354
+ iou_thresh: float,
355
+ ) -> np.ndarray:
356
+ """
357
+ Standard NMS: keep one box per overlapping cluster (the one with highest score).
358
+ Returns indices of kept boxes (into the boxes/scores arrays).
359
+ """
360
+ N = len(boxes)
361
+ if N == 0:
362
+ return np.array([], dtype=np.intp)
363
+ boxes = np.asarray(boxes, dtype=np.float32)
364
+ scores = np.asarray(scores, dtype=np.float32)
365
+ order = np.argsort(scores)[::-1]
366
+ keep: list[int] = []
367
+ suppressed = np.zeros(N, dtype=bool)
368
+ for i in range(N):
369
+ idx = order[i]
370
+ if suppressed[idx]:
371
+ continue
372
+ keep.append(idx)
373
+ bi = boxes[idx]
374
+ for k in range(i + 1, N):
375
+ jdx = order[k]
376
+ if suppressed[jdx]:
377
+ continue
378
+ bj = boxes[jdx]
379
+ xx1 = max(bi[0], bj[0])
380
+ yy1 = max(bi[1], bj[1])
381
+ xx2 = min(bi[2], bj[2])
382
+ yy2 = min(bi[3], bj[3])
383
+ inter = max(0.0, xx2 - xx1) * max(0.0, yy2 - yy1)
384
+ area_i = (bi[2] - bi[0]) * (bi[3] - bi[1])
385
+ area_j = (bj[2] - bj[0]) * (bj[3] - bj[1])
386
+ iou = inter / (area_i + area_j - inter + 1e-7)
387
+ if iou > iou_thresh:
388
+ suppressed[jdx] = True
389
+ return np.array(keep)
390
+
391
+ def _per_class_hard_nms(
392
+ self,
393
+ boxes: np.ndarray,
394
+ scores: np.ndarray,
395
+ cls_ids: np.ndarray,
396
+ iou_thresh: float,
397
+ ) -> np.ndarray:
398
+ """Hard NMS applied independently per class."""
399
+ if len(boxes) == 0:
400
+ return np.array([], dtype=np.intp)
401
+ all_keep: list[int] = []
402
+ for c in np.unique(cls_ids):
403
+ mask = cls_ids == c
404
+ indices = np.where(mask)[0]
405
+ keep = self._hard_nms(boxes[mask], scores[mask], iou_thresh)
406
+ all_keep.extend(indices[keep].tolist())
407
+ all_keep.sort()
408
+ return np.array(all_keep, dtype=np.intp)
409
+
410
+ def _per_class_soft_nms(
411
+ self,
412
+ boxes: np.ndarray,
413
+ scores: np.ndarray,
414
+ cls_ids: np.ndarray,
415
+ sigma: float = 0.5,
416
+ score_thresh: float = 0.01,
417
+ ) -> tuple[np.ndarray, np.ndarray]:
418
+ """Soft NMS applied independently per class."""
419
+ if len(boxes) == 0:
420
+ return np.array([], dtype=np.intp), np.array([], dtype=np.float32)
421
+ all_keep: list[int] = []
422
+ all_scores: list[float] = []
423
+ for c in np.unique(cls_ids):
424
+ mask = cls_ids == c
425
+ indices = np.where(mask)[0]
426
+ keep, updated = self._soft_nms(boxes[mask], scores[mask], sigma, score_thresh)
427
+ for k, s in zip(keep, updated):
428
+ all_keep.append(int(indices[k]))
429
+ all_scores.append(float(s))
430
+ if not all_keep:
431
+ return np.array([], dtype=np.intp), np.array([], dtype=np.float32)
432
+ return np.array(all_keep, dtype=np.intp), np.array(all_scores, dtype=np.float32)
433
+
434
+ def _filter_sane_boxes(
435
+ self,
436
+ boxes: np.ndarray,
437
+ scores: np.ndarray,
438
+ cls_ids: np.ndarray,
439
+ orig_size: tuple[int, int],
440
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
441
+ """Filter out tiny, degenerate, or implausible boxes (common FP)."""
442
+ if len(boxes) == 0:
443
+ return boxes, scores, cls_ids
444
+ orig_w, orig_h = orig_size
445
+ image_area = float(orig_w * orig_h)
446
+ keep = []
447
+ for i, box in enumerate(boxes):
448
+ x1, y1, x2, y2 = box.tolist()
449
+ bw = x2 - x1
450
+ bh = y2 - y1
451
+ if bw <= 0 or bh <= 0:
452
+ continue
453
+ if bw < self.min_side or bh < self.min_side:
454
+ continue
455
+ area = bw * bh
456
+ if area < self.min_box_area:
457
+ continue
458
+ if area > 0.95 * image_area:
459
+ continue
460
+ ar = max(bw / max(bh, 1e-6), bh / max(bw, 1e-6))
461
+ if ar > self.max_aspect_ratio:
462
+ continue
463
+ keep.append(i)
464
+ if not keep:
465
+ return (
466
+ np.empty((0, 4), dtype=np.float32),
467
+ np.empty((0,), dtype=np.float32),
468
+ np.empty((0,), dtype=np.int32),
469
+ )
470
+ k = np.array(keep, dtype=np.intp)
471
+ return boxes[k], scores[k], cls_ids[k]
472
+
473
+ @staticmethod
474
+ def _max_score_per_cluster(
475
+ post_boxes: np.ndarray,
476
+ post_cls: np.ndarray,
477
+ full_boxes: np.ndarray,
478
+ full_scores: np.ndarray,
479
+ full_cls: np.ndarray,
480
+ iou_thresh: float,
481
+ ) -> np.ndarray:
482
+ """For each kept (post-NMS) box, return the max score over the FULL
483
+ candidate set among SAME-CLASS boxes with IoU >= iou_thresh.
484
+
485
+ The previous version omitted the same-class constraint, which let a
486
+ confident broom raise the score of a coincident nozzle (or vice
487
+ versa) under TTA. That's a silent FP booster and is fixed here.
488
+ """
489
+ n = len(post_boxes)
490
+ if n == 0:
491
+ return np.empty(0, dtype=np.float32)
492
+ full_areas = (np.maximum(0.0, full_boxes[:, 2] - full_boxes[:, 0]) *
493
+ np.maximum(0.0, full_boxes[:, 3] - full_boxes[:, 1]))
494
+ out = np.empty(n, dtype=np.float32)
495
+ for i in range(n):
496
+ bi = post_boxes[i]
497
+ xx1 = np.maximum(bi[0], full_boxes[:, 0])
498
+ yy1 = np.maximum(bi[1], full_boxes[:, 1])
499
+ xx2 = np.minimum(bi[2], full_boxes[:, 2])
500
+ yy2 = np.minimum(bi[3], full_boxes[:, 3])
501
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
502
+ a_i = max(0.0, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
503
+ iou = inter / (a_i + full_areas - inter + 1e-7)
504
+ cluster = (iou >= iou_thresh) & (full_cls == post_cls[i])
505
+ out[i] = float(np.max(full_scores[cluster])) if np.any(cluster) else 0.0
506
+ return out
507
+
508
+ def _conf_filter_mask(
509
+ self, scores: np.ndarray, cls_ids: np.ndarray
510
+ ) -> np.ndarray:
511
+ """Boolean keep-mask: score >= per-class threshold, with a per-class
512
+ rescue -- if a class has zero boxes passing, admit its top-1 candidate
513
+ when its score >= (per-class threshold - per-class bonus).
514
+ """
515
+ if len(scores) == 0:
516
+ return np.zeros(0, dtype=bool)
517
+ thr = self._conf_thres_array[cls_ids]
518
+ keep = scores >= thr
519
+ for c in np.unique(cls_ids):
520
+ b = float(self._bonus_array[c])
521
+ if b <= 0.0:
522
+ continue
523
+ cm = cls_ids == c
524
+ if keep[cm].any():
525
+ continue
526
+ idx = np.where(cm)[0]
527
+ top = int(idx[int(np.argmax(scores[idx]))])
528
+ if scores[top] >= self._conf_thres_array[c] - b:
529
+ keep[top] = True
530
+ return keep
531
+
532
+ def _cross_class_dedup_op(
533
+ self,
534
+ boxes: np.ndarray,
535
+ scores: np.ndarray,
536
+ cls_ids: np.ndarray,
537
+ iou_thresh: float,
538
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
539
+ """Remove near-duplicate boxes across classes.
540
+
541
+ Order candidates by (score - per_class_threshold) margin, then by area;
542
+ keep the highest, suppress every other box with IoU > iou_thresh. For
543
+ car-wash this kills the common failure where water spray makes the
544
+ model fire both `nozzle` and `track` on the same patch, or where a
545
+ broom handle overlaps a drainage-gate detection.
546
+ """
547
+ n = len(boxes)
548
+ if n <= 1:
549
+ return boxes, scores, cls_ids
550
+ boxes = np.asarray(boxes, dtype=np.float32)
551
+ scores = np.asarray(scores, dtype=np.float32)
552
+ cls_ids = np.asarray(cls_ids, dtype=np.int32)
553
+ areas = (np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) *
554
+ np.maximum(0.0, boxes[:, 3] - boxes[:, 1]))
555
+ margins = scores - self._conf_thres_array[cls_ids]
556
+ order = np.lexsort((-areas, -margins))
557
+ suppressed = np.zeros(n, dtype=bool)
558
+ keep: list[int] = []
559
+ for i in order:
560
+ if suppressed[i]:
561
+ continue
562
+ keep.append(int(i))
563
+ bi = boxes[i]
564
+ xx1 = np.maximum(bi[0], boxes[:, 0])
565
+ yy1 = np.maximum(bi[1], boxes[:, 1])
566
+ xx2 = np.minimum(bi[2], boxes[:, 2])
567
+ yy2 = np.minimum(bi[3], boxes[:, 3])
568
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
569
+ a_i = max(1e-7, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
570
+ iou = inter / (a_i + areas - inter + 1e-7)
571
+ dup = iou > iou_thresh
572
+ dup[i] = False
573
+ suppressed |= dup
574
+ keep_idx = np.array(keep, dtype=np.intp)
575
+ return boxes[keep_idx], scores[keep_idx], cls_ids[keep_idx]
576
+
577
+ def _per_view_pipeline(
578
+ self,
579
+ boxes: np.ndarray,
580
+ scores: np.ndarray,
581
+ cls_ids: np.ndarray,
582
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
583
+ """Per-view post-processing: per-class NMS -> cap -> cross-class dedup."""
584
+ if len(boxes) > 1:
585
+ keep = self._per_class_hard_nms(boxes, scores, cls_ids, self.iou_thres)
586
+ boxes, scores, cls_ids = boxes[keep], scores[keep], cls_ids[keep]
587
+ if len(scores) > self.max_det:
588
+ top = np.argsort(-scores)[: self.max_det]
589
+ boxes, scores, cls_ids = boxes[top], scores[top], cls_ids[top]
590
+ if len(boxes) > 1:
591
+ boxes, scores, cls_ids = self._cross_class_dedup_op(
592
+ boxes, scores, cls_ids, self.cross_iou_thresh
593
+ )
594
+ return boxes, scores, cls_ids
595
+
596
+ def _decode_final_dets(
597
+ self,
598
+ preds: np.ndarray,
599
+ ratio: float,
600
+ pad: tuple[float, float],
601
+ orig_size: tuple[int, int],
602
+ apply_optional_dedup: bool = False,
603
+ ) -> list[BoundingBox]:
604
+ """
605
+ Primary path:
606
+ expected output rows like [x1, y1, x2, y2, conf, cls_id]
607
+ in letterboxed input coordinates.
608
+ """
609
+ if preds.ndim == 3 and preds.shape[0] == 1:
610
+ preds = preds[0]
611
+
612
+ if preds.ndim != 2 or preds.shape[1] < 6:
613
+ raise ValueError(f"Unexpected ONNX final-det output shape: {preds.shape}")
614
+
615
+ boxes = preds[:, :4].astype(np.float32)
616
+ scores = preds[:, 4].astype(np.float32)
617
+ cls_ids = preds[:, 5].astype(np.int32)
618
+ cls_ids = self.cls_remap[cls_ids]
619
+
620
+ # Per-class confidence filter with rescue (replaces scalar threshold)
621
+ keep = self._conf_filter_mask(scores, cls_ids)
622
+ boxes = boxes[keep]
623
+ scores = scores[keep]
624
+ cls_ids = cls_ids[keep]
625
+
626
+ if len(boxes) == 0:
627
+ return []
628
+
629
+ pad_w, pad_h = pad
630
+ orig_w, orig_h = orig_size
631
+
632
+ # reverse letterbox
633
+ boxes[:, [0, 2]] -= pad_w
634
+ boxes[:, [1, 3]] -= pad_h
635
+ boxes /= ratio
636
+ boxes = self._clip_boxes(boxes, (orig_w, orig_h))
637
+
638
+ # Box sanity filter (reduces FP)
639
+ boxes, scores, cls_ids = self._filter_sane_boxes(
640
+ boxes, scores, cls_ids, orig_size
641
+ )
642
+ if len(boxes) == 0:
643
+ return []
644
+
645
+ if apply_optional_dedup and len(boxes) > 1:
646
+ # Soft-NMS path preserved as a tunable option; default below.
647
+ keep_idx, scores = self._per_class_soft_nms(boxes, scores, cls_ids)
648
+ boxes = boxes[keep_idx]
649
+ cls_ids = cls_ids[keep_idx]
650
+ if len(scores) > self.max_det:
651
+ top = np.argsort(-scores)[: self.max_det]
652
+ boxes, scores, cls_ids = boxes[top], scores[top], cls_ids[top]
653
+ if len(boxes) > 1:
654
+ boxes, scores, cls_ids = self._cross_class_dedup_op(
655
+ boxes, scores, cls_ids, self.cross_iou_thresh
656
+ )
657
+ else:
658
+ # Default: per-class hard NMS -> cap -> cross-class dedup
659
+ boxes, scores, cls_ids = self._per_view_pipeline(boxes, scores, cls_ids)
660
+
661
+ results: list[BoundingBox] = []
662
+ for box, conf, cls_id in zip(boxes, scores, cls_ids):
663
+ x1, y1, x2, y2 = box.tolist()
664
+
665
+ if x2 <= x1 or y2 <= y1:
666
+ continue
667
+
668
+ results.append(
669
+ BoundingBox(
670
+ x1=int(math.floor(x1)),
671
+ y1=int(math.floor(y1)),
672
+ x2=int(math.ceil(x2)),
673
+ y2=int(math.ceil(y2)),
674
+ cls_id=int(cls_id),
675
+ conf=float(conf),
676
+ )
677
+ )
678
+
679
+ return results
680
+
681
+ def _decode_raw_yolo(
682
+ self,
683
+ preds: np.ndarray,
684
+ ratio: float,
685
+ pad: tuple[float, float],
686
+ orig_size: tuple[int, int],
687
+ ) -> list[BoundingBox]:
688
+ """
689
+ Fallback path for raw YOLO predictions.
690
+ Supports common layouts:
691
+ - [1, C, N]
692
+ - [1, N, C]
693
+ """
694
+ if preds.ndim != 3:
695
+ raise ValueError(f"Unexpected raw ONNX output shape: {preds.shape}")
696
+
697
+ if preds.shape[0] != 1:
698
+ raise ValueError(f"Unexpected batch dimension in raw output: {preds.shape}")
699
+
700
+ preds = preds[0]
701
+
702
+ # Normalize to [N, C]
703
+ if preds.shape[0] <= 16 and preds.shape[1] > preds.shape[0]:
704
+ preds = preds.T
705
+
706
+ if preds.ndim != 2 or preds.shape[1] < 5:
707
+ raise ValueError(f"Unexpected normalized raw output shape: {preds.shape}")
708
+
709
+ boxes_xywh = preds[:, :4].astype(np.float32)
710
+ cls_part = preds[:, 4:].astype(np.float32)
711
+
712
+ if cls_part.shape[1] == 1:
713
+ scores = cls_part[:, 0]
714
+ cls_ids = np.zeros(len(scores), dtype=np.int32)
715
+ else:
716
+ cls_ids = np.argmax(cls_part, axis=1).astype(np.int32)
717
+ scores = cls_part[np.arange(len(cls_part)), cls_ids]
718
+ cls_ids = self.cls_remap[cls_ids]
719
+
720
+ # Per-class confidence filter with rescue (replaces scalar threshold)
721
+ keep = self._conf_filter_mask(scores, cls_ids)
722
+ boxes_xywh = boxes_xywh[keep]
723
+ scores = scores[keep]
724
+ cls_ids = cls_ids[keep]
725
+ if len(boxes_xywh) == 0:
726
+ return []
727
+
728
+ boxes = self._xywh_to_xyxy(boxes_xywh)
729
+
730
+ # Order matches fire001 / _decode_final_dets:
731
+ # unscale -> clip -> sanity filter -> per-view pipeline (NMS, cap, cross-class dedup).
732
+ pad_w, pad_h = pad
733
+ orig_w, orig_h = orig_size
734
+ boxes[:, [0, 2]] -= pad_w
735
+ boxes[:, [1, 3]] -= pad_h
736
+ boxes /= ratio
737
+ boxes = self._clip_boxes(boxes, (orig_w, orig_h))
738
+
739
+ boxes, scores, cls_ids = self._filter_sane_boxes(
740
+ boxes, scores, cls_ids, (orig_w, orig_h)
741
+ )
742
+ if len(boxes) == 0:
743
+ return []
744
+
745
+ boxes, scores, cls_ids = self._per_view_pipeline(boxes, scores, cls_ids)
746
+
747
+ results: list[BoundingBox] = []
748
+ for box, conf, cls_id in zip(boxes, scores, cls_ids):
749
+ x1, y1, x2, y2 = box.tolist()
750
+
751
+ if x2 <= x1 or y2 <= y1:
752
+ continue
753
+
754
+ results.append(
755
+ BoundingBox(
756
+ x1=int(math.floor(x1)),
757
+ y1=int(math.floor(y1)),
758
+ x2=int(math.ceil(x2)),
759
+ y2=int(math.ceil(y2)),
760
+ cls_id=int(cls_id),
761
+ conf=float(conf),
762
+ )
763
+ )
764
+
765
+ return results
766
+
767
+ def _postprocess(
768
+ self,
769
+ output: np.ndarray,
770
+ ratio: float,
771
+ pad: tuple[float, float],
772
+ orig_size: tuple[int, int],
773
+ ) -> list[BoundingBox]:
774
+ """
775
+ Prefer final detections first.
776
+ Fallback to raw decode only if needed.
777
+ """
778
+ # final detections: [N,6]
779
+ if output.ndim == 2 and output.shape[1] >= 6:
780
+ return self._decode_final_dets(output, ratio, pad, orig_size)
781
+
782
+ # final detections: [1,N,6]
783
+ if output.ndim == 3 and output.shape[0] == 1 and output.shape[2] == 6:
784
+ return self._decode_final_dets(output, ratio, pad, orig_size)
785
+
786
+ # fallback raw decode
787
+ return self._decode_raw_yolo(output, ratio, pad, orig_size)
788
+
789
+ def _predict_single(self, image: np.ndarray) -> list[BoundingBox]:
790
+ if image is None:
791
+ raise ValueError("Input image is None")
792
+ if not isinstance(image, np.ndarray):
793
+ raise TypeError(f"Input is not numpy array: {type(image)}")
794
+ if image.ndim != 3:
795
+ raise ValueError(f"Expected HWC image, got shape={image.shape}")
796
+ if image.shape[0] <= 0 or image.shape[1] <= 0:
797
+ raise ValueError(f"Invalid image shape={image.shape}")
798
+ if image.shape[2] != 3:
799
+ raise ValueError(f"Expected 3 channels, got shape={image.shape}")
800
+
801
+ if image.dtype != np.uint8:
802
+ image = image.astype(np.uint8)
803
+
804
+ input_tensor, ratio, pad, orig_size = self._preprocess(image)
805
+
806
+ expected_shape = (1, 3, self.input_height, self.input_width)
807
+ if input_tensor.shape != expected_shape:
808
+ raise ValueError(
809
+ f"Bad input tensor shape={input_tensor.shape}, expected={expected_shape}"
810
+ )
811
+
812
+ outputs = self.session.run(self.output_names, {self.input_name: input_tensor})
813
+ det_output = outputs[0]
814
+ return self._postprocess(det_output, ratio, pad, orig_size)
815
+
816
+ def _predict_tta(self, image: np.ndarray) -> list[BoundingBox]:
817
+ """Horizontal-flip TTA.
818
+
819
+ Strategy (ported from fire001):
820
+ 1. Predict on original and on flipped image.
821
+ 2. Map flipped boxes back to original coordinates.
822
+ 3. Per-class hard NMS on the union.
823
+ 4. For each kept box, compute the max SAME-CLASS score across the
824
+ FULL union -- a high-confidence flipped detection raises a
825
+ borderline original one, but never one of a different class.
826
+ 5. Cross-class dedup to suppress same-physical-object multi-class.
827
+ """
828
+ boxes_orig = self._predict_single(image)
829
+
830
+ flipped = cv2.flip(image, 1)
831
+ boxes_flip = self._predict_single(flipped)
832
+
833
+ w = image.shape[1]
834
+ boxes_flip = [
835
+ BoundingBox(
836
+ x1=w - b.x2, y1=b.y1, x2=w - b.x1, y2=b.y2,
837
+ cls_id=b.cls_id, conf=b.conf,
838
+ )
839
+ for b in boxes_flip
840
+ ]
841
+
842
+ all_boxes = boxes_orig + boxes_flip
843
+ if len(all_boxes) == 0:
844
+ return []
845
+
846
+ coords = np.array(
847
+ [[b.x1, b.y1, b.x2, b.y2] for b in all_boxes], dtype=np.float32
848
+ )
849
+ scores = np.array([b.conf for b in all_boxes], dtype=np.float32)
850
+ cls_ids = np.array([b.cls_id for b in all_boxes], dtype=np.int32)
851
+
852
+ hard_keep = self._per_class_hard_nms(coords, scores, cls_ids, self.iou_thres)
853
+ if len(hard_keep) == 0:
854
+ return []
855
+ if len(hard_keep) > self.max_det:
856
+ top = np.argsort(-scores[hard_keep])[: self.max_det]
857
+ hard_keep = hard_keep[top]
858
+
859
+ # Class-aware cluster-max score boost (fixes the silent cross-class
860
+ # leak in the previous _max_score_per_cluster).
861
+ boosted = self._max_score_per_cluster(
862
+ coords[hard_keep], cls_ids[hard_keep],
863
+ coords, scores, cls_ids, self.iou_thres,
864
+ )
865
+
866
+ kept_coords = coords[hard_keep]
867
+ kept_cls = cls_ids[hard_keep]
868
+ if len(kept_coords) > 1:
869
+ kept_coords, boosted, kept_cls = self._cross_class_dedup_op(
870
+ kept_coords, boosted, kept_cls, self.cross_iou_thresh
871
+ )
872
+
873
+ return [
874
+ BoundingBox(
875
+ x1=int(math.floor(kept_coords[j, 0])),
876
+ y1=int(math.floor(kept_coords[j, 1])),
877
+ x2=int(math.ceil(kept_coords[j, 2])),
878
+ y2=int(math.ceil(kept_coords[j, 3])),
879
+ cls_id=int(kept_cls[j]),
880
+ conf=float(boosted[j]),
881
+ )
882
+ for j in range(len(kept_coords))
883
+ ]
884
+
885
+ def predict_batch(
886
+ self,
887
+ batch_images: list[ndarray],
888
+ offset: int,
889
+ n_keypoints: int,
890
+ ) -> list[TVFrameResult]:
891
+ results: list[TVFrameResult] = []
892
+
893
+ for frame_number_in_batch, image in enumerate(batch_images):
894
+ try:
895
+ if self.use_tta:
896
+ boxes = self._predict_tta(image)
897
+ else:
898
+ boxes = self._predict_single(image)
899
+ except Exception as e:
900
+ print(f"⚠️ Inference failed for frame {offset + frame_number_in_batch}: {e}")
901
+ boxes = []
902
+
903
+ results.append(
904
+ TVFrameResult(
905
+ frame_id=offset + frame_number_in_batch,
906
+ boxes=boxes,
907
+ keypoints=[(0, 0) for _ in range(max(0, int(n_keypoints)))],
908
+ )
909
+ )
910
+
911
+ return results
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
weights.onnx CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:5eb4543ec44fe7dd6a743e36b687ddd08985e54fdd358426b118bec1c4972ad7
3
- size 9805054
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:0e6e1480df66e695ec12ca70116ff30edff54de3473fdd95c0c64ef902e6a9be
3
+ size 9842746