SuperBitDev commited on
Commit
95b1912
Β·
verified Β·
1 Parent(s): 20a4942

Upload folder using huggingface_hub

Browse files
Files changed (3) hide show
  1. class_names.txt +4 -0
  2. miner.py +420 -684
  3. weights.onnx +2 -2
class_names.txt ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ broom
2
+ drainage gate
3
+ nozzle
4
+ track
miner.py CHANGED
@@ -24,249 +24,141 @@ class TVFrameResult(BaseModel):
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
-
54
- try:
55
- self.session = ort.InferenceSession(
56
- str(model_path),
57
- sess_options=sess_options,
58
- providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
59
- )
60
- print("βœ… Created ORT session with preferred CUDA provider list")
61
- except Exception as e:
62
- print(f"⚠️ CUDA session creation failed, falling back to CPU: {e}")
63
- self.session = ort.InferenceSession(
64
- str(model_path),
65
- sess_options=sess_options,
66
- providers=["CPUExecutionProvider"],
67
- )
68
-
69
- print("ORT session providers:", self.session.get_providers())
70
-
71
- # Build cls_remap: for each model-emit index i,
72
- # cls_remap[i] = self.class_names.index(model_class_order[i])
73
- # The model-side order comes from the ONNX metadata when available,
74
- # else falls back to the static _model_class_order.
75
- model_class_order = self._read_model_class_order()
76
- if model_class_order is None:
77
- model_class_order = list(self._model_class_order)
78
- print(f"cls order: no usable ONNX metadata, FALLBACK {model_class_order}")
79
- else:
80
- print(f"cls order: from ONNX metadata {model_class_order}")
81
- self.cls_remap = np.array(
82
- [self.class_names.index(n) for n in model_class_order], dtype=np.int32
83
  )
84
-
85
- for inp in self.session.get_inputs():
86
- print("INPUT:", inp.name, inp.shape, inp.type)
87
-
88
- for out in self.session.get_outputs():
89
- print("OUTPUT:", out.name, out.shape, out.type)
90
-
91
  self.input_name = self.session.get_inputs()[0].name
92
  self.output_names = [output.name for output in self.session.get_outputs()]
93
- self.input_shape = self.session.get_inputs()[0].shape
94
-
95
- # Match the ONNX input dtype (this export is FP16 -> needs float16 input).
96
- input_type = self.session.get_inputs()[0].type
97
- self.np_dtype = np.float16 if "float16" in input_type else np.float32
98
- print(f"βœ… ONNX input dtype: {input_type} -> numpy {self.np_dtype}")
99
-
100
- # ONNX is fixed-size 1408x1408 (v1 export); read actual shape to be safe.
101
- self.input_height = self._safe_dim(self.input_shape[2], default=1280)
102
- self.input_width = self._safe_dim(self.input_shape[3], default=1280)
103
-
104
- # Tuned for validator scoring (pillars: 0.6*map50 + 0.4*false_positive).
105
- # All values below are the measured optimum of a full grid sweep on
106
- # the validator-style val split (tune_miner.py, 241 1024x1024 crops,
107
- # composite 0.8002 -> 0.8103) -- re-run the sweep after any retrain.
108
- self.iou_thres = 0.45 # Per-class NMS IoU; lower = stricter dedup
109
- self.cross_iou_thresh = 0.9 # Cross-class dedup IoU (suppress same physical object firing multiple classes)
110
- self.max_det = 200
111
- self.use_tta = True
112
-
113
- # conf thresholds: broom=0.38 drainage gate=0.45 nozzle=0.30 track=0.60
114
- # Per-class confidence thresholds.
115
- # Indexed by class_names order: [broom, drainage gate, nozzle, track].
116
- # broom/nozzle sit low: under the validator metric the mAP gained
117
- # from the extra recall outweighs the FP-pillar cost (the previous
118
- # 0.5/0.5 silently discarded many valid detections); track is the
119
- # one class where false fires are common enough to need 0.38.
120
- self._conf_thres_array = np.array(
121
- [0.37, 0.23, 0.37, 0.45], dtype=np.float32
122
- )
123
- # Per-class rescue bonus: when a class has ZERO boxes passing the
124
- # threshold in a frame, its top-1 candidate is admitted when its score
125
- # is at least (per-class threshold - per-class bonus).
126
- # DISABLED (all zeros): the sweep showed rescue admits more false
127
- # positives than true positives under the validator's FP pillar.
128
- self._bonus_array = np.array(
129
- [0.05, 0.05, 0.0, 0.15], dtype=np.float32
130
- )
131
-
132
- # Box sanity filter β€” kept loose: car-wash `nozzle` boxes are tiny
133
- # (GT median ~290 pxΒ², smallest ~32 pxΒ²). Fire's 14x14/min_side 8
134
- # would delete valid nozzles, so thresholds are dropped here.
135
- self.min_box_area = 4 * 4 # 16 pxΒ²
136
- self.min_side = 3
137
- self.max_aspect_ratio = 12.0
138
-
139
- print(f"βœ… ONNX model loaded from: {model_path}")
140
- print(f"βœ… ONNX providers: {self.session.get_providers()}")
141
- print(f"βœ… ONNX input: name={self.input_name}, shape={self.input_shape}")
142
 
143
  def __repr__(self) -> str:
144
- return (
145
- f"ONNXRuntime(session={type(self.session).__name__}, "
146
- f"providers={self.session.get_providers()})"
147
- )
148
 
149
  @staticmethod
150
  def _safe_dim(value, default: int) -> int:
151
  return value if isinstance(value, int) and value > 0 else default
152
 
153
- @staticmethod
154
- def _resolve_model_path(repo: Path) -> Path:
155
- """Locate the ONNX model in the repo dir.
156
-
157
- Prefers weights.onnx (FP16/FP32 export), then weights_int8.onnx (the
158
- training script's INT8-quantized export -- works as-is: quantization
159
- preserves the Ultralytics metadata and QDQ models take regular fp32
160
- input), then any other .onnx file. INT8 is the fallback when the FP16
161
- export exceeds the 30 MB deployment limit (e.g. yolo26m).
162
- """
163
- for name in ("weights.onnx", "weights_int8.onnx"):
164
- p = repo / name
165
- if p.exists():
166
- if name != "weights.onnx":
167
- print(f"model: weights.onnx not found, using {name}")
168
- return p
169
- candidates = sorted(repo.glob("*.onnx"))
170
- if candidates:
171
- print(f"model: using {candidates[0].name}")
172
- return candidates[0]
173
- return repo / "weights.onnx" # let session creation raise the error
174
-
175
- def _read_model_class_order(self) -> list[str] | None:
176
- """Read the model's class order from Ultralytics ONNX metadata.
177
-
178
- Returns the class names ordered by model-emit index, or None when
179
- metadata is missing/unparsable or doesn't match `class_names` as a
180
- set (in which case the static _model_class_order fallback is used).
181
- """
182
- try:
183
- import ast
184
-
185
- meta = self.session.get_modelmeta().custom_metadata_map
186
- names = ast.literal_eval(meta["names"]) # e.g. {0: 'broom', ...}
187
- if isinstance(names, dict):
188
- order = [str(names[i]) for i in sorted(names)]
189
- else:
190
- order = [str(n) for n in names]
191
- except Exception as e:
192
- print(f"cls order: could not read ONNX names metadata ({e})")
193
- return None
194
- if sorted(order) != sorted(self.class_names):
195
- print(
196
- f"cls order: ONNX names {order} do not match expected classes "
197
- f"{self.class_names}; ignoring metadata"
198
- )
199
- return None
200
- return order
201
-
202
  def _letterbox(
203
- self,
204
- image: ndarray,
205
- new_shape: tuple[int, int],
206
- color=(114, 114, 114),
207
- ) -> tuple[ndarray, float, tuple[float, float]]:
208
- """
209
- Resize with unchanged aspect ratio and pad to target shape.
210
- Returns:
211
- padded_image,
212
- ratio,
213
- (pad_w, pad_h) # half-padding
214
- """
215
- h, w = image.shape[:2]
216
- new_w, new_h = new_shape
217
-
218
- ratio = min(new_w / w, new_h / h)
219
- resized_w = int(round(w * ratio))
220
- resized_h = int(round(h * ratio))
221
-
222
- if (resized_w, resized_h) != (w, h):
223
- interp = cv2.INTER_CUBIC if ratio > 1.0 else cv2.INTER_LINEAR
224
- image = cv2.resize(image, (resized_w, resized_h), interpolation=interp)
225
-
226
- dw = new_w - resized_w
227
- dh = new_h - resized_h
228
- dw /= 2.0
229
- dh /= 2.0
230
-
231
- left = int(round(dw - 0.1))
232
- right = int(round(dw + 0.1))
233
- top = int(round(dh - 0.1))
234
- bottom = int(round(dh + 0.1))
235
-
236
- padded = cv2.copyMakeBorder(
237
- image,
238
- top,
239
- bottom,
240
- left,
241
- right,
242
- borderType=cv2.BORDER_CONSTANT,
243
- value=color,
244
- )
245
- return padded, ratio, (dw, dh)
246
-
247
- def _preprocess(
248
- self, image: ndarray
249
- ) -> tuple[np.ndarray, float, tuple[float, float], tuple[int, int]]:
250
- """
251
- Preprocess for fixed-size ONNX export:
252
- - enhance image quality (CLAHE, denoise, sharpen)
253
- - letterbox to model input size
254
- - BGR -> RGB
255
- - normalize to [0,1]
256
- - HWC -> NCHW float32
257
- """
258
  orig_h, orig_w = image.shape[:2]
259
-
260
- img, ratio, pad = self._letterbox(
261
- image, (self.input_width, self.input_height)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
262
  )
263
- img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
264
- img = (img.astype(np.float32) / 255.0)
265
- img = np.transpose(img, (2, 0, 1))[None, ...]
266
- img = np.ascontiguousarray(img, dtype=self.np_dtype)
267
-
268
- return img, ratio, pad, (orig_w, orig_h)
269
-
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
270
  @staticmethod
271
  def _clip_boxes(boxes: np.ndarray, image_size: tuple[int, int]) -> np.ndarray:
272
  w, h = image_size
@@ -277,200 +169,116 @@ class Miner:
277
  return boxes
278
 
279
  @staticmethod
280
- def _xywh_to_xyxy(boxes: np.ndarray) -> np.ndarray:
281
- out = np.empty_like(boxes)
282
- out[:, 0] = boxes[:, 0] - boxes[:, 2] / 2.0
283
- out[:, 1] = boxes[:, 1] - boxes[:, 3] / 2.0
284
- out[:, 2] = boxes[:, 0] + boxes[:, 2] / 2.0
285
- out[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.0
286
- return out
287
-
288
- def _soft_nms(
289
- self,
290
- boxes: np.ndarray,
291
- scores: np.ndarray,
292
- sigma: float = 0.5,
293
- score_thresh: float = 0.01,
294
- ) -> tuple[np.ndarray, np.ndarray]:
295
- """
296
- Soft-NMS: Gaussian decay of overlapping scores instead of hard removal.
297
- Returns (kept_original_indices, updated_scores).
298
- """
299
- N = len(boxes)
300
- if N == 0:
301
- return np.array([], dtype=np.intp), np.array([], dtype=np.float32)
302
-
303
- boxes = boxes.astype(np.float32, copy=True)
304
- scores = scores.astype(np.float32, copy=True)
305
- order = np.arange(N)
306
-
307
- for i in range(N):
308
- max_pos = i + int(np.argmax(scores[i:]))
309
- boxes[[i, max_pos]] = boxes[[max_pos, i]]
310
- scores[[i, max_pos]] = scores[[max_pos, i]]
311
- order[[i, max_pos]] = order[[max_pos, i]]
312
-
313
- if i + 1 >= N:
314
  break
315
-
316
- xx1 = np.maximum(boxes[i, 0], boxes[i + 1:, 0])
317
- yy1 = np.maximum(boxes[i, 1], boxes[i + 1:, 1])
318
- xx2 = np.minimum(boxes[i, 2], boxes[i + 1:, 2])
319
- yy2 = np.minimum(boxes[i, 3], boxes[i + 1:, 3])
320
  inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
321
-
322
- area_i = max(0.0, float(
323
- (boxes[i, 2] - boxes[i, 0]) * (boxes[i, 3] - boxes[i, 1])
324
- ))
325
- areas_j = (
326
- np.maximum(0.0, boxes[i + 1:, 2] - boxes[i + 1:, 0])
327
- * np.maximum(0.0, boxes[i + 1:, 3] - boxes[i + 1:, 1])
328
- )
329
- iou = inter / (area_i + areas_j - inter + 1e-7)
330
- scores[i + 1:] *= np.exp(-(iou ** 2) / sigma)
331
-
332
- mask = scores > score_thresh
333
- return order[mask], scores[mask]
334
-
335
- @staticmethod
336
- def _hard_nms(
337
- boxes: np.ndarray,
338
- scores: np.ndarray,
339
- iou_thresh: float,
340
- ) -> np.ndarray:
341
- """
342
- Standard NMS: keep one box per overlapping cluster (the one with highest score).
343
- Returns indices of kept boxes (into the boxes/scores arrays).
344
- """
345
- N = len(boxes)
346
- if N == 0:
347
- return np.array([], dtype=np.intp)
348
- boxes = np.asarray(boxes, dtype=np.float32)
349
- scores = np.asarray(scores, dtype=np.float32)
350
- order = np.argsort(scores)[::-1]
351
- keep: list[int] = []
352
- suppressed = np.zeros(N, dtype=bool)
353
- for i in range(N):
354
- idx = order[i]
355
- if suppressed[idx]:
356
- continue
357
- keep.append(idx)
358
- bi = boxes[idx]
359
- for k in range(i + 1, N):
360
- jdx = order[k]
361
- if suppressed[jdx]:
362
- continue
363
- bj = boxes[jdx]
364
- xx1 = max(bi[0], bj[0])
365
- yy1 = max(bi[1], bj[1])
366
- xx2 = min(bi[2], bj[2])
367
- yy2 = min(bi[3], bj[3])
368
- inter = max(0.0, xx2 - xx1) * max(0.0, yy2 - yy1)
369
- area_i = (bi[2] - bi[0]) * (bi[3] - bi[1])
370
- area_j = (bj[2] - bj[0]) * (bj[3] - bj[1])
371
- iou = inter / (area_i + area_j - inter + 1e-7)
372
- if iou > iou_thresh:
373
- suppressed[jdx] = True
374
- return np.array(keep)
375
-
376
- def _per_class_hard_nms(
377
- self,
378
- boxes: np.ndarray,
379
- scores: np.ndarray,
380
- cls_ids: np.ndarray,
381
- iou_thresh: float,
382
- ) -> np.ndarray:
383
- """Hard NMS applied independently per class."""
384
  if len(boxes) == 0:
385
  return np.array([], dtype=np.intp)
386
- all_keep: list[int] = []
387
  for c in np.unique(cls_ids):
388
  mask = cls_ids == c
389
  indices = np.where(mask)[0]
390
- keep = self._hard_nms(boxes[mask], scores[mask], iou_thresh)
 
391
  all_keep.extend(indices[keep].tolist())
392
  all_keep.sort()
393
  return np.array(all_keep, dtype=np.intp)
394
 
395
- def _per_class_soft_nms(
396
- self,
397
- boxes: np.ndarray,
398
- scores: np.ndarray,
399
- cls_ids: np.ndarray,
400
- sigma: float = 0.5,
401
- score_thresh: float = 0.01,
402
- ) -> tuple[np.ndarray, np.ndarray]:
403
- """Soft NMS applied independently per class."""
404
- if len(boxes) == 0:
405
- return np.array([], dtype=np.intp), np.array([], dtype=np.float32)
406
- all_keep: list[int] = []
407
- all_scores: list[float] = []
408
- for c in np.unique(cls_ids):
409
- mask = cls_ids == c
410
- indices = np.where(mask)[0]
411
- keep, updated = self._soft_nms(boxes[mask], scores[mask], sigma, score_thresh)
412
- for k, s in zip(keep, updated):
413
- all_keep.append(int(indices[k]))
414
- all_scores.append(float(s))
415
- if not all_keep:
416
- return np.array([], dtype=np.intp), np.array([], dtype=np.float32)
417
- return np.array(all_keep, dtype=np.intp), np.array(all_scores, dtype=np.float32)
418
-
419
- def _filter_sane_boxes(
420
- self,
421
- boxes: np.ndarray,
422
- scores: np.ndarray,
423
- cls_ids: np.ndarray,
424
- orig_size: tuple[int, int],
425
- ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
426
- """Filter out tiny, degenerate, or implausible boxes (common FP)."""
427
- if len(boxes) == 0:
428
  return boxes, scores, cls_ids
429
- orig_w, orig_h = orig_size
430
- image_area = float(orig_w * orig_h)
 
 
 
 
 
 
431
  keep = []
432
- for i, box in enumerate(boxes):
433
- x1, y1, x2, y2 = box.tolist()
434
- bw = x2 - x1
435
- bh = y2 - y1
436
- if bw <= 0 or bh <= 0:
437
- continue
438
- if bw < self.min_side or bh < self.min_side:
439
- continue
440
- area = bw * bh
441
- if area < self.min_box_area:
442
- continue
443
- if area > 0.95 * image_area:
444
- continue
445
- ar = max(bw / max(bh, 1e-6), bh / max(bw, 1e-6))
446
- if ar > self.max_aspect_ratio:
447
  continue
448
- keep.append(i)
449
- if not keep:
450
- return (
451
- np.empty((0, 4), dtype=np.float32),
452
- np.empty((0,), dtype=np.float32),
453
- np.empty((0,), dtype=np.int32),
454
- )
455
- k = np.array(keep, dtype=np.intp)
456
- return boxes[k], scores[k], cls_ids[k]
 
 
 
 
 
457
 
458
- @staticmethod
459
- def _max_score_per_cluster(
460
- post_boxes: np.ndarray,
461
- post_cls: np.ndarray,
462
- full_boxes: np.ndarray,
463
- full_scores: np.ndarray,
464
- full_cls: np.ndarray,
465
- iou_thresh: float,
466
- ) -> np.ndarray:
467
- """For each kept (post-NMS) box, return the max score over the FULL
468
- candidate set among SAME-CLASS boxes with IoU >= iou_thresh.
469
-
470
- The previous version omitted the same-class constraint, which let a
471
- confident broom raise the score of a coincident nozzle (or vice
472
- versa) under TTA. That's a silent FP booster and is fixed here.
473
- """
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
474
  n = len(post_boxes)
475
  if n == 0:
476
  return np.empty(0, dtype=np.float32)
@@ -490,16 +298,16 @@ class Miner:
490
  out[i] = float(np.max(full_scores[cluster])) if np.any(cluster) else 0.0
491
  return out
492
 
493
- def _conf_filter_mask(
494
- self, scores: np.ndarray, cls_ids: np.ndarray
495
- ) -> np.ndarray:
496
- """Boolean keep-mask: score >= per-class threshold, with a per-class
497
- rescue -- if a class has zero boxes passing, admit its top-1 candidate
498
- when its score >= (per-class threshold - per-class bonus).
499
- """
500
  if len(scores) == 0:
501
  return np.zeros(0, dtype=bool)
502
- thr = self._conf_thres_array[cls_ids]
 
 
 
 
503
  keep = scores >= thr
504
  for c in np.unique(cls_ids):
505
  b = float(self._bonus_array[c])
@@ -514,60 +322,82 @@ class Miner:
514
  keep[top] = True
515
  return keep
516
 
517
- def _cross_class_dedup_op(
518
  self,
519
  boxes: np.ndarray,
520
  scores: np.ndarray,
521
  cls_ids: np.ndarray,
522
- iou_thresh: float,
523
  ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
524
- """Remove near-duplicate boxes across classes.
 
 
 
 
 
 
525
 
526
- Order candidates by (score - per_class_threshold) margin, then by area;
527
- keep the highest, suppress every other box with IoU > iou_thresh. For
528
- car-wash this kills the common failure where water spray makes the
529
- model fire both `nozzle` and `track` on the same patch, or where a
530
- broom handle overlaps a drainage-gate detection.
531
  """
532
  n = len(boxes)
533
  if n <= 1:
534
  return boxes, scores, cls_ids
 
535
  boxes = np.asarray(boxes, dtype=np.float32)
536
- scores = np.asarray(scores, dtype=np.float32)
537
  cls_ids = np.asarray(cls_ids, dtype=np.int32)
 
538
  areas = (np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) *
539
- np.maximum(0.0, boxes[:, 3] - boxes[:, 1]))
540
- margins = scores - self._conf_thres_array[cls_ids]
541
- order = np.lexsort((-areas, -margins))
542
- suppressed = np.zeros(n, dtype=bool)
543
- keep: list[int] = []
544
- for i in order:
545
- if suppressed[i]:
 
 
 
546
  continue
547
- keep.append(int(i))
548
- bi = boxes[i]
549
- xx1 = np.maximum(bi[0], boxes[:, 0])
550
- yy1 = np.maximum(bi[1], boxes[:, 1])
551
- xx2 = np.minimum(bi[2], boxes[:, 2])
552
- yy2 = np.minimum(bi[3], boxes[:, 3])
553
- inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
554
- a_i = max(1e-7, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
555
- iou = inter / (a_i + areas - inter + 1e-7)
556
- dup = iou > iou_thresh
557
- dup[i] = False
558
- suppressed |= dup
559
- keep_idx = np.array(keep, dtype=np.intp)
 
 
 
 
 
 
 
 
 
 
560
  return boxes[keep_idx], scores[keep_idx], cls_ids[keep_idx]
561
 
562
- def _per_view_pipeline(
563
- self,
564
- boxes: np.ndarray,
565
- scores: np.ndarray,
566
- cls_ids: np.ndarray,
567
- ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
568
- """Per-view post-processing: per-class NMS -> cap -> cross-class dedup."""
 
 
569
  if len(boxes) > 1:
570
- keep = self._per_class_hard_nms(boxes, scores, cls_ids, self.iou_thres)
571
  boxes, scores, cls_ids = boxes[keep], scores[keep], cls_ids[keep]
572
  if len(scores) > self.max_det:
573
  top = np.argsort(-scores)[: self.max_det]
@@ -578,295 +408,206 @@ class Miner:
578
  )
579
  return boxes, scores, cls_ids
580
 
581
- def _decode_final_dets(
582
- self,
583
- preds: np.ndarray,
584
- ratio: float,
585
- pad: tuple[float, float],
586
- orig_size: tuple[int, int],
587
- apply_optional_dedup: bool = False,
588
- ) -> list[BoundingBox]:
589
- """
590
- Primary path:
591
- expected output rows like [x1, y1, x2, y2, conf, cls_id]
592
- in letterboxed input coordinates.
593
- """
594
  if preds.ndim == 3 and preds.shape[0] == 1:
595
  preds = preds[0]
596
 
597
  if preds.ndim != 2 or preds.shape[1] < 6:
598
- raise ValueError(f"Unexpected ONNX final-det output shape: {preds.shape}")
 
599
 
600
  boxes = preds[:, :4].astype(np.float32)
601
  scores = preds[:, 4].astype(np.float32)
602
  cls_ids = preds[:, 5].astype(np.int32)
603
- cls_ids = self.cls_remap[cls_ids]
604
 
605
- # Per-class confidence filter with rescue (replaces scalar threshold)
606
- keep = self._conf_filter_mask(scores, cls_ids)
607
- boxes = boxes[keep]
608
- scores = scores[keep]
609
- cls_ids = cls_ids[keep]
610
 
611
  if len(boxes) == 0:
612
  return []
613
 
614
- pad_w, pad_h = pad
615
- orig_w, orig_h = orig_size
616
-
617
- # reverse letterbox
618
- boxes[:, [0, 2]] -= pad_w
619
- boxes[:, [1, 3]] -= pad_h
620
- boxes /= ratio
621
- boxes = self._clip_boxes(boxes, (orig_w, orig_h))
622
 
623
- # Box sanity filter (reduces FP)
624
- boxes, scores, cls_ids = self._filter_sane_boxes(
625
- boxes, scores, cls_ids, orig_size
626
- )
627
- if len(boxes) == 0:
628
- return []
629
-
630
- if apply_optional_dedup and len(boxes) > 1:
631
- # Soft-NMS path preserved as a tunable option; default below.
632
- keep_idx, scores = self._per_class_soft_nms(boxes, scores, cls_ids)
633
- boxes = boxes[keep_idx]
634
- cls_ids = cls_ids[keep_idx]
635
- if len(scores) > self.max_det:
636
- top = np.argsort(-scores)[: self.max_det]
637
- boxes, scores, cls_ids = boxes[top], scores[top], cls_ids[top]
638
- if len(boxes) > 1:
639
- boxes, scores, cls_ids = self._cross_class_dedup_op(
640
- boxes, scores, cls_ids, self.cross_iou_thresh
641
- )
642
- else:
643
- # Default: per-class hard NMS -> cap -> cross-class dedup
644
- boxes, scores, cls_ids = self._per_view_pipeline(boxes, scores, cls_ids)
645
-
646
- results: list[BoundingBox] = []
647
- for box, conf, cls_id in zip(boxes, scores, cls_ids):
648
- x1, y1, x2, y2 = box.tolist()
649
-
650
- if x2 <= x1 or y2 <= y1:
651
- continue
652
-
653
- results.append(
654
- BoundingBox(
655
- x1=int(math.floor(x1)),
656
- y1=int(math.floor(y1)),
657
- x2=int(math.ceil(x2)),
658
- y2=int(math.ceil(y2)),
659
- cls_id=int(cls_id),
660
- conf=float(conf),
661
- )
662
- )
663
-
664
- return results
665
-
666
- def _decode_raw_yolo(
667
- self,
668
- preds: np.ndarray,
669
- ratio: float,
670
- pad: tuple[float, float],
671
- orig_size: tuple[int, int],
672
- ) -> list[BoundingBox]:
673
- """
674
- Fallback path for raw YOLO predictions.
675
- Supports common layouts:
676
- - [1, C, N]
677
- - [1, N, C]
678
- """
679
- if preds.ndim != 3:
680
- raise ValueError(f"Unexpected raw ONNX output shape: {preds.shape}")
681
-
682
- if preds.shape[0] != 1:
683
- raise ValueError(f"Unexpected batch dimension in raw output: {preds.shape}")
684
-
685
- preds = preds[0]
686
-
687
- # Normalize to [N, C]
688
- if preds.shape[0] <= 16 and preds.shape[1] > preds.shape[0]:
689
- preds = preds.T
690
-
691
- if preds.ndim != 2 or preds.shape[1] < 5:
692
- raise ValueError(f"Unexpected normalized raw output shape: {preds.shape}")
693
-
694
- boxes_xywh = preds[:, :4].astype(np.float32)
695
- cls_part = preds[:, 4:].astype(np.float32)
696
-
697
- if cls_part.shape[1] == 1:
698
- scores = cls_part[:, 0]
699
- cls_ids = np.zeros(len(scores), dtype=np.int32)
700
- else:
701
- cls_ids = np.argmax(cls_part, axis=1).astype(np.int32)
702
- scores = cls_part[np.arange(len(cls_part)), cls_ids]
703
- cls_ids = self.cls_remap[cls_ids]
704
-
705
- # Per-class confidence filter with rescue (replaces scalar threshold)
706
- keep = self._conf_filter_mask(scores, cls_ids)
707
- boxes_xywh = boxes_xywh[keep]
708
  scores = scores[keep]
709
  cls_ids = cls_ids[keep]
710
- if len(boxes_xywh) == 0:
711
- return []
712
 
713
- boxes = self._xywh_to_xyxy(boxes_xywh)
 
714
 
715
- # Order matches fire001 / _decode_final_dets:
716
- # unscale -> clip -> sanity filter -> per-view pipeline (NMS, cap, cross-class dedup).
717
  pad_w, pad_h = pad
718
- orig_w, orig_h = orig_size
719
  boxes[:, [0, 2]] -= pad_w
720
  boxes[:, [1, 3]] -= pad_h
721
  boxes /= ratio
722
- boxes = self._clip_boxes(boxes, (orig_w, orig_h))
723
 
724
- boxes, scores, cls_ids = self._filter_sane_boxes(
725
- boxes, scores, cls_ids, (orig_w, orig_h)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
726
  )
727
- if len(boxes) == 0:
728
- return []
729
 
730
- boxes, scores, cls_ids = self._per_view_pipeline(boxes, scores, cls_ids)
731
 
732
- results: list[BoundingBox] = []
 
 
 
 
 
733
  for box, conf, cls_id in zip(boxes, scores, cls_ids):
734
  x1, y1, x2, y2 = box.tolist()
735
-
736
  if x2 <= x1 or y2 <= y1:
737
  continue
738
-
739
  results.append(
740
  BoundingBox(
741
- x1=int(math.floor(x1)),
742
- y1=int(math.floor(y1)),
743
- x2=int(math.ceil(x2)),
744
- y2=int(math.ceil(y2)),
745
  cls_id=int(cls_id),
746
- conf=float(conf),
747
  )
748
  )
749
-
750
  return results
751
 
752
- def _postprocess(
753
- self,
754
- output: np.ndarray,
755
- ratio: float,
756
- pad: tuple[float, float],
757
- orig_size: tuple[int, int],
758
- ) -> list[BoundingBox]:
759
- """
760
- Prefer final detections first.
761
- Fallback to raw decode only if needed.
762
- """
763
- # final detections: [N,6]
764
- if output.ndim == 2 and output.shape[1] >= 6:
765
- return self._decode_final_dets(output, ratio, pad, orig_size)
766
-
767
- # final detections: [1,N,6]
768
- if output.ndim == 3 and output.shape[0] == 1 and output.shape[2] == 6:
769
- return self._decode_final_dets(output, ratio, pad, orig_size)
770
-
771
- # fallback raw decode
772
- return self._decode_raw_yolo(output, ratio, pad, orig_size)
773
-
774
- def _predict_single(self, image: np.ndarray) -> list[BoundingBox]:
775
- if image is None:
776
- raise ValueError("Input image is None")
777
- if not isinstance(image, np.ndarray):
778
- raise TypeError(f"Input is not numpy array: {type(image)}")
779
- if image.ndim != 3:
780
- raise ValueError(f"Expected HWC image, got shape={image.shape}")
781
- if image.shape[0] <= 0 or image.shape[1] <= 0:
782
- raise ValueError(f"Invalid image shape={image.shape}")
783
- if image.shape[2] != 3:
784
- raise ValueError(f"Expected 3 channels, got shape={image.shape}")
785
-
786
- if image.dtype != np.uint8:
787
- image = image.astype(np.uint8)
788
-
789
- input_tensor, ratio, pad, orig_size = self._preprocess(image)
790
-
791
- expected_shape = (1, 3, self.input_height, self.input_width)
792
- if input_tensor.shape != expected_shape:
793
- raise ValueError(
794
- f"Bad input tensor shape={input_tensor.shape}, expected={expected_shape}"
795
- )
796
-
797
- outputs = self.session.run(self.output_names, {self.input_name: input_tensor})
798
- det_output = outputs[0]
799
- return self._postprocess(det_output, ratio, pad, orig_size)
800
-
801
- def _predict_tta(self, image: np.ndarray) -> list[BoundingBox]:
802
- """Horizontal-flip TTA.
803
-
804
- Strategy (ported from fire001):
805
- 1. Predict on original and on flipped image.
806
- 2. Map flipped boxes back to original coordinates.
807
- 3. Per-class hard NMS on the union.
808
- 4. For each kept box, compute the max SAME-CLASS score across the
809
- FULL union -- a high-confidence flipped detection raises a
810
- borderline original one, but never one of a different class.
811
- 5. Cross-class dedup to suppress same-physical-object multi-class.
812
- """
813
- boxes_orig = self._predict_single(image)
814
-
815
- flipped = cv2.flip(image, 1)
816
- boxes_flip = self._predict_single(flipped)
817
-
818
- w = image.shape[1]
819
  boxes_flip = [
820
- BoundingBox(
821
- x1=w - b.x2, y1=b.y1, x2=w - b.x1, y2=b.y2,
822
- cls_id=b.cls_id, conf=b.conf,
823
- )
824
  for b in boxes_flip
825
  ]
826
-
827
- all_boxes = boxes_orig + boxes_flip
828
- if len(all_boxes) == 0:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
829
  return []
830
 
831
- coords = np.array(
832
- [[b.x1, b.y1, b.x2, b.y2] for b in all_boxes], dtype=np.float32
833
- )
834
  scores = np.array([b.conf for b in all_boxes], dtype=np.float32)
835
  cls_ids = np.array([b.cls_id for b in all_boxes], dtype=np.int32)
836
 
837
- hard_keep = self._per_class_hard_nms(coords, scores, cls_ids, self.iou_thres)
 
838
  if len(hard_keep) == 0:
839
  return []
 
840
  if len(hard_keep) > self.max_det:
841
  top = np.argsort(-scores[hard_keep])[: self.max_det]
842
  hard_keep = hard_keep[top]
843
-
844
- # Class-aware cluster-max score boost (fixes the silent cross-class
845
- # leak in the previous _max_score_per_cluster).
846
  boosted = self._max_score_per_cluster(
847
  coords[hard_keep], cls_ids[hard_keep],
848
- coords, scores, cls_ids, self.iou_thres,
849
  )
850
 
851
  kept_coords = coords[hard_keep]
852
  kept_cls = cls_ids[hard_keep]
 
 
853
  if len(kept_coords) > 1:
854
  kept_coords, boosted, kept_cls = self._cross_class_dedup_op(
855
  kept_coords, boosted, kept_cls, self.cross_iou_thresh
856
  )
857
 
858
- return [
859
- BoundingBox(
860
- x1=int(math.floor(kept_coords[j, 0])),
861
- y1=int(math.floor(kept_coords[j, 1])),
862
- x2=int(math.ceil(kept_coords[j, 2])),
863
- y2=int(math.ceil(kept_coords[j, 3])),
864
- cls_id=int(kept_cls[j]),
865
- conf=float(boosted[j]),
 
 
 
 
 
 
866
  )
867
- for j in range(len(kept_coords))
868
- ]
869
 
 
 
870
  def predict_batch(
871
  self,
872
  batch_images: list[ndarray],
@@ -874,23 +615,18 @@ class Miner:
874
  n_keypoints: int,
875
  ) -> list[TVFrameResult]:
876
  results: list[TVFrameResult] = []
877
-
878
- for frame_number_in_batch, image in enumerate(batch_images):
879
  try:
880
- if self.use_tta:
881
- boxes = self._predict_tta(image)
882
- else:
883
- boxes = self._predict_single(image)
884
  except Exception as e:
885
- print(f"⚠️ Inference failed for frame {offset + frame_number_in_batch}: {e}")
886
  boxes = []
887
-
888
  results.append(
889
  TVFrameResult(
890
- frame_id=offset + frame_number_in_batch,
891
  boxes=boxes,
892
- keypoints=[(0, 0) for _ in range(max(0, int(n_keypoints)))],
893
  )
894
  )
895
-
896
  return results
 
24
 
25
 
26
  class Miner:
27
+ """
28
+ YOLOv26 ONNX miner for car wash detection.
29
+
30
+ Classes: broom, drainage gate, nozzle, track
31
+
32
+ v26 is NMS-free β€” output shape: [1, 300, 6] (x1, y1, x2, y2, conf, cls_id).
33
+
34
+ Features:
35
+ - Vectorized NMS + sanity filter + dedup + flip TTA
36
+ - Per-class rescue bonus (saves hard-to-detect classes at slightly lower conf)
37
+ - Confidence boost from same-class cluster (TTA consensus)
38
+ - Aggressive same-class overlap suppression
39
+ - Per-class IoU thresholds
40
+ """
41
+
42
+ class_names = ['broom', 'drainage gate', 'nozzle', 'track']
43
+ input_size = 1408
44
+ cross_iou_thresh = 0.8
45
+ max_det = 300
46
+ #overlap_suppress_threshold = 0.85
47
+
48
+ # Per-class confidence thresholds
49
+ _conf_thres_array = np.array([0.35, 0.7, 0.4, 0.7], dtype=np.float32)
50
+ _extra_conf_thres_array = np.array([0.35, 0.35, 0.55, 0.35], dtype=np.float32)
51
+
52
+ # Per-class IoU thresholds for same-class NMS
53
+ _iou_thres_array = np.array([0.6, 0.65, 0.5, 0.65], dtype=np.float32)
54
+
55
+ # Per-class rescue bonus
56
+ _bonus_array = np.array([0.2, 0.2, 0.0, 0.2], dtype=np.float32)
57
+
58
+ # Per-class minimum box area
59
+ # Indices: 0=broom, 1=drainage gate, 2=nozzle, 3=track
60
+ _min_box_area_array = np.array([144.0, 144.0, 4.0, 144.0], dtype=np.float32)
61
+
62
+
63
+ def __init__(self, path_hf_repo: Path) -> None:
64
+ self.path_hf_repo = path_hf_repo
65
+
66
  print("ORT version:", ort.__version__)
67
+
68
  try:
69
  ort.preload_dlls()
70
+ print("preload_dlls success")
71
  except Exception as e:
72
+ print(f"preload_dlls failed: {e}")
73
+
74
  print("ORT available providers BEFORE session:", ort.get_available_providers())
75
+
76
  sess_options = ort.SessionOptions()
77
  sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
78
+
79
+ self.session = ort.InferenceSession(
80
+ str(path_hf_repo / "weights.onnx"),
81
+ sess_options=sess_options,
82
+ providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
83
  )
84
+ print("Created ORT session with preferred CUDA provider list")
85
+ print("ORT session providers:", self.session.get_providers())
86
+
 
 
 
 
87
  self.input_name = self.session.get_inputs()[0].name
88
  self.output_names = [output.name for output in self.session.get_outputs()]
89
+ input_shape = self.session.get_inputs()[0].shape
90
+
91
+ self.input_h = self._safe_dim(input_shape[2], default=self.input_size)
92
+ self.input_w = self._safe_dim(input_shape[3], default=self.input_size)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
93
 
94
  def __repr__(self) -> str:
95
+ return f"YOLOv26 Car Wash Miner classes={len(self.class_names)}"
 
 
 
96
 
97
  @staticmethod
98
  def _safe_dim(value, default: int) -> int:
99
  return value if isinstance(value, int) and value > 0 else default
100
 
101
+ # ─── Preprocessing ────────────────────────────────────────────
102
+
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
103
  def _letterbox(
104
+ self, image: ndarray, new_shape: tuple[int, int],
105
+ color: tuple[int, int, int] = (114, 114, 114),
106
+ ) -> tuple[ndarray, float, float, float]:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
107
  orig_h, orig_w = image.shape[:2]
108
+ target_w, target_h = new_shape
109
+
110
+ r = min(target_w / orig_w, target_h / orig_h)
111
+ new_unpad_w = int(round(orig_w * r))
112
+ new_unpad_h = int(round(orig_h * r))
113
+
114
+ resized = cv2.resize(image, (new_unpad_w, new_unpad_h), interpolation=cv2.INTER_LINEAR)
115
+
116
+ dw = target_w - new_unpad_w
117
+ dh = target_h - new_unpad_h
118
+ pad_w = dw / 2.0
119
+ pad_h = dh / 2.0
120
+
121
+ left = int(round(pad_w - 0.1))
122
+ right = int(round(pad_w + 0.1))
123
+ top = int(round(pad_h - 0.1))
124
+ bottom = int(round(pad_h + 0.1))
125
+
126
+ out = cv2.copyMakeBorder(
127
+ resized, top, bottom, left, right,
128
+ cv2.BORDER_CONSTANT, value=color,
129
  )
130
+ return out, r, pad_w, pad_h
131
+
132
+ def _preprocess(self, image_bgr: np.ndarray,
133
+ allow_pad: bool = True) -> tuple[np.ndarray, dict]:
134
+ orig_h, orig_w = image_bgr.shape[:2]
135
+ extra_left = 0
136
+ extra_right = 0
137
+ if allow_pad and orig_w == orig_h: # only pad when allowed
138
+ target_w = int(orig_w * 1.05)
139
+ if target_w > orig_w:
140
+ total_extra = target_w - orig_w
141
+ extra_left = total_extra // 2
142
+ extra_right = total_extra - extra_left
143
+ image_bgr = cv2.copyMakeBorder(
144
+ image_bgr, 0, 0, extra_left, extra_right,
145
+ cv2.BORDER_CONSTANT, value=(114, 114, 114),
146
+ )
147
+ padded_h, padded_w = image_bgr.shape[:2]
148
+ rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
149
+ img, ratio, pad_w, pad_h = self._letterbox(rgb, (self.input_w, self.input_h))
150
+ x = img.astype(np.float32) / 255.0
151
+ x = np.transpose(x, (2, 0, 1))[None, ...]
152
+ x = np.ascontiguousarray(x)
153
+ return x, {
154
+ "orig_h": orig_h, "orig_w": orig_w,
155
+ "ratio": ratio, "pad_w": pad_w, "pad_h": pad_h,
156
+ "extra_left": extra_left, "extra_right": extra_right,
157
+ "padded_w": padded_w, "padded_h": padded_h,
158
+ }
159
+
160
+ # ─── Vectorized box operations ───────────────────────────────
161
+
162
  @staticmethod
163
  def _clip_boxes(boxes: np.ndarray, image_size: tuple[int, int]) -> np.ndarray:
164
  w, h = image_size
 
169
  return boxes
170
 
171
  @staticmethod
172
+ def _hard_nms(boxes: np.ndarray, scores: np.ndarray,
173
+ iou_thresh: float) -> np.ndarray:
174
+ """Vectorized NMS. Returns indices to keep."""
175
+ n = len(boxes)
176
+ if n == 0:
177
+ return np.array([], dtype=np.intp)
178
+ order = np.argsort(-scores)
179
+ keep = []
180
+ while len(order) > 0:
181
+ i = int(order[0])
182
+ keep.append(i)
183
+ if len(order) == 1:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
184
  break
185
+ rest = order[1:]
186
+ xx1 = np.maximum(boxes[i, 0], boxes[rest, 0])
187
+ yy1 = np.maximum(boxes[i, 1], boxes[rest, 1])
188
+ xx2 = np.minimum(boxes[i, 2], boxes[rest, 2])
189
+ yy2 = np.minimum(boxes[i, 3], boxes[rest, 3])
190
  inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
191
+ a_i = (max(0.0, boxes[i, 2] - boxes[i, 0]) *
192
+ max(0.0, boxes[i, 3] - boxes[i, 1]))
193
+ a_r = (np.maximum(0.0, boxes[rest, 2] - boxes[rest, 0]) *
194
+ np.maximum(0.0, boxes[rest, 3] - boxes[rest, 1]))
195
+ iou = inter / (a_i + a_r - inter + 1e-7)
196
+ order = rest[iou <= iou_thresh]
197
+ return np.array(keep, dtype=np.intp)
198
+
199
+ def _per_class_hard_nms(self, boxes: np.ndarray, scores: np.ndarray,
200
+ cls_ids: np.ndarray) -> np.ndarray:
201
+ """Per-class NMS using per-class IoU thresholds."""
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
202
  if len(boxes) == 0:
203
  return np.array([], dtype=np.intp)
204
+ all_keep = []
205
  for c in np.unique(cls_ids):
206
  mask = cls_ids == c
207
  indices = np.where(mask)[0]
208
+ cls_iou = float(self._iou_thres_array[c]) # per-class IoU threshold
209
+ keep = self._hard_nms(boxes[mask], scores[mask], cls_iou)
210
  all_keep.extend(indices[keep].tolist())
211
  all_keep.sort()
212
  return np.array(all_keep, dtype=np.intp)
213
 
214
+ def _cross_class_dedup_op(self, boxes: np.ndarray, scores: np.ndarray,
215
+ cls_ids: np.ndarray, iou_thresh: float
216
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
217
+ n = len(boxes)
218
+ if n <= 1:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
219
  return boxes, scores, cls_ids
220
+ boxes = np.asarray(boxes, dtype=np.float32)
221
+ scores = np.asarray(scores, dtype=np.float32)
222
+ cls_ids = np.asarray(cls_ids, dtype=np.int32)
223
+ areas = (np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) *
224
+ np.maximum(0.0, boxes[:, 3] - boxes[:, 1]))
225
+ margins = scores - self._conf_thres_array[cls_ids]
226
+ order = np.lexsort((-areas, -margins))
227
+ suppressed = np.zeros(n, dtype=bool)
228
  keep = []
229
+ for i in order:
230
+ if suppressed[i]:
 
 
 
 
 
 
 
 
 
 
 
 
 
231
  continue
232
+ keep.append(int(i))
233
+ bi = boxes[i]
234
+ xx1 = np.maximum(bi[0], boxes[:, 0])
235
+ yy1 = np.maximum(bi[1], boxes[:, 1])
236
+ xx2 = np.minimum(bi[2], boxes[:, 2])
237
+ yy2 = np.minimum(bi[3], boxes[:, 3])
238
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
239
+ a_i = max(1e-7, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
240
+ iou = inter / (a_i + areas - inter + 1e-7)
241
+ dup = iou > iou_thresh
242
+ dup[i] = False
243
+ suppressed |= dup
244
+ keep_idx = np.array(keep, dtype=np.intp)
245
+ return boxes[keep_idx], scores[keep_idx], cls_ids[keep_idx]
246
 
247
+ def _filter_sane_boxes(self, boxes: np.ndarray, scores: np.ndarray,
248
+ cls_ids: np.ndarray, orig_size: tuple[int, int]
249
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
250
+ """Filter by per-class min area, max area ratio, and aspect ratio."""
251
+ if len(boxes) == 0:
252
+ return boxes, scores, cls_ids
253
+
254
+ orig_w, orig_h = orig_size
255
+ image_area = float(orig_w * orig_h)
256
+ bw = np.maximum(0.0, boxes[:, 2] - boxes[:, 0])
257
+ bh = np.maximum(0.0, boxes[:, 3] - boxes[:, 1])
258
+ area = bw * bh
259
+
260
+ ar = np.where(
261
+ (bw > 0) & (bh > 0),
262
+ np.maximum(bw / np.maximum(bh, 1e-6), bh / np.maximum(bw, 1e-6)),
263
+ np.inf,
264
+ )
265
+
266
+ # Per-class minimum area
267
+ class_min_area = self._min_box_area_array[cls_ids]
268
+
269
+ keep = (
270
+ (area >= class_min_area) &
271
+ (area <= 0.95 * image_area)
272
+ )
273
+ return boxes[keep], scores[keep], cls_ids[keep]
274
+
275
+ def _max_score_per_cluster(self, post_boxes: np.ndarray,
276
+ post_cls: np.ndarray,
277
+ full_boxes: np.ndarray,
278
+ full_scores: np.ndarray,
279
+ full_cls: np.ndarray,
280
+ iou_thresh: float) -> np.ndarray:
281
+ """For each kept box, set confidence to max score in its SAME-CLASS cluster."""
282
  n = len(post_boxes)
283
  if n == 0:
284
  return np.empty(0, dtype=np.float32)
 
298
  out[i] = float(np.max(full_scores[cluster])) if np.any(cluster) else 0.0
299
  return out
300
 
301
+ def _conf_filter_mask(self, scores: np.ndarray,
302
+ cls_ids: np.ndarray, extra_left: int) -> np.ndarray:
303
+ """Per-class threshold with rescue bonus for missed classes."""
 
 
 
 
304
  if len(scores) == 0:
305
  return np.zeros(0, dtype=bool)
306
+ thr = 0
307
+ if extra_left > 0:
308
+ thr = self._extra_conf_thres_array[cls_ids]
309
+ else:
310
+ thr = self._conf_thres_array[cls_ids]
311
  keep = scores >= thr
312
  for c in np.unique(cls_ids):
313
  b = float(self._bonus_array[c])
 
322
  keep[top] = True
323
  return keep
324
 
325
+ def _suppress_overlapping_same_class(
326
  self,
327
  boxes: np.ndarray,
328
  scores: np.ndarray,
329
  cls_ids: np.ndarray,
330
+ threshold: float,
331
  ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
332
+ """
333
+ Drop a same-class box that is (almost) entirely *contained* inside a larger
334
+ same-class box β€” a duplicate detection of one object.
335
+
336
+ Containment is intersection / area_of_SMALLER_box (IoMin), NOT IoU.
337
+ A small box nested in a large one has tiny IoU, so plain NMS never removes
338
+ it; IoMin catches it.
339
 
340
+ black (large) + green (fully inside black) -> same gate, drop green
341
+ red (sticks out of black) -> separate gate, keep
342
+
343
+ Survivor = the LARGER box, and its confidence is raised to the cluster max.
 
344
  """
345
  n = len(boxes)
346
  if n <= 1:
347
  return boxes, scores, cls_ids
348
+
349
  boxes = np.asarray(boxes, dtype=np.float32)
350
+ scores = np.asarray(scores, dtype=np.float32).copy()
351
  cls_ids = np.asarray(cls_ids, dtype=np.int32)
352
+
353
  areas = (np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) *
354
+ np.maximum(0.0, boxes[:, 3] - boxes[:, 1]))
355
+
356
+ keep = np.ones(n, dtype=bool)
357
+
358
+ # Largest first, so the survivor of a containment chain is the biggest box.
359
+ order = np.argsort(-areas)
360
+
361
+ for idx_a in range(n):
362
+ a = order[idx_a]
363
+ if not keep[a]:
364
  continue
365
+ for idx_b in range(idx_a + 1, n):
366
+ b = order[idx_b] # areas[b] <= areas[a]
367
+ if not keep[b]:
368
+ continue
369
+ if cls_ids[a] != cls_ids[b]:
370
+ continue
371
+
372
+ x1 = max(boxes[a, 0], boxes[b, 0])
373
+ y1 = max(boxes[a, 1], boxes[b, 1])
374
+ x2 = min(boxes[a, 2], boxes[b, 2])
375
+ y2 = min(boxes[a, 3], boxes[b, 3])
376
+ if x2 <= x1 or y2 <= y1:
377
+ continue
378
+ inter = (x2 - x1) * (y2 - y1)
379
+
380
+ # How much of the SMALLER box (b) lies inside the larger (a):
381
+ containment_b = inter / max(areas[b], 1e-9)
382
+
383
+ if containment_b >= threshold: # b is nested -> it's a duplicate
384
+ scores[a] = max(scores[a], scores[b]) # keep the higher score
385
+ keep[b] = False # drop the smaller (green)
386
+
387
+ keep_idx = np.where(keep)[0]
388
  return boxes[keep_idx], scores[keep_idx], cls_ids[keep_idx]
389
 
390
+ def _per_view_pipeline(self, boxes: np.ndarray, scores: np.ndarray,
391
+ cls_ids: np.ndarray, orig_size: tuple[int, int]
392
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
393
+ """Sanity filter + per-class NMS + cross-class dedup."""
394
+ boxes, scores, cls_ids = self._filter_sane_boxes(
395
+ boxes, scores, cls_ids, orig_size
396
+ )
397
+ if len(boxes) == 0:
398
+ return boxes, scores, cls_ids
399
  if len(boxes) > 1:
400
+ keep = self._per_class_hard_nms(boxes, scores, cls_ids)
401
  boxes, scores, cls_ids = boxes[keep], scores[keep], cls_ids[keep]
402
  if len(scores) > self.max_det:
403
  top = np.argsort(-scores)[: self.max_det]
 
408
  )
409
  return boxes, scores, cls_ids
410
 
411
+ # ─── v26-specific decoding ────────────────────────────────────
412
+
413
+ def _decode_v26_output(self, preds: np.ndarray, ratio: float,
414
+ pad: tuple[float, float],
415
+ orig_size: tuple[int, int],
416
+ extra: tuple[int, int] = (0, 0)) -> list[BoundingBox]: # NEW arg
417
+ """Decode YOLOv26 output (shape [1, 300, 6] or [300, 6])."""
 
 
 
 
 
 
418
  if preds.ndim == 3 and preds.shape[0] == 1:
419
  preds = preds[0]
420
 
421
  if preds.ndim != 2 or preds.shape[1] < 6:
422
+ print(f"Warning: Unexpected v26 output shape: {preds.shape}")
423
+ return []
424
 
425
  boxes = preds[:, :4].astype(np.float32)
426
  scores = preds[:, 4].astype(np.float32)
427
  cls_ids = preds[:, 5].astype(np.int32)
 
428
 
429
+ n_cls = len(self.class_names)
430
+ valid = (cls_ids >= 0) & (cls_ids < n_cls)
431
+ boxes = boxes[valid]
432
+ scores = scores[valid]
433
+ cls_ids = cls_ids[valid]
434
 
435
  if len(boxes) == 0:
436
  return []
437
 
438
+ extra_left, _extra_right = extra
 
 
 
 
 
 
 
439
 
440
+ keep = self._conf_filter_mask(scores, cls_ids, extra_left)
441
+ boxes = boxes[keep]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
442
  scores = scores[keep]
443
  cls_ids = cls_ids[keep]
 
 
444
 
445
+ if len(boxes) == 0:
446
+ return []
447
 
448
+ # 1) undo letterbox -> coords in the PADDED (widened) image
 
449
  pad_w, pad_h = pad
 
450
  boxes[:, [0, 2]] -= pad_w
451
  boxes[:, [1, 3]] -= pad_h
452
  boxes /= ratio
 
453
 
454
+ # 2) NEW: undo the left/right pre-padding -> coords in the ORIGINAL image.
455
+ # Only the LEFT pad shifts x; right pad adds width but no offset.
456
+
457
+ if extra_left:
458
+ boxes[:, [0, 2]] -= extra_left
459
+
460
+ # 2b) NEW: drop boxes that fall in the black padding region.
461
+ # A real detection must have its CENTER inside the original image
462
+ # width [0, orig_w]; boxes centered in the black bars are spurious.
463
+ if extra_left or _extra_right:
464
+ orig_w, orig_h = orig_size
465
+ cx = (boxes[:, 0] + boxes[:, 2]) * 0.5
466
+ inside = (cx >= 0) & (cx <= orig_w)
467
+ boxes = boxes[inside]
468
+ scores = scores[inside]
469
+ cls_ids = cls_ids[inside]
470
+ if len(boxes) == 0:
471
+ return []
472
+
473
+ # 3) clip to ORIGINAL image bounds (orig_size is the true original size)
474
+ boxes = self._clip_boxes(boxes, orig_size)
475
+
476
+ boxes, scores, cls_ids = self._per_view_pipeline(
477
+ boxes, scores, cls_ids, orig_size
478
  )
 
 
479
 
480
+ return self._build_results(boxes, scores, cls_ids, orig_size)
481
 
482
+ @staticmethod
483
+ def _build_results(boxes: np.ndarray, scores: np.ndarray,
484
+ cls_ids: np.ndarray,
485
+ orig_size: tuple[int, int]) -> list[BoundingBox]:
486
+ results = []
487
+ orig_w, orig_h = orig_size
488
  for box, conf, cls_id in zip(boxes, scores, cls_ids):
489
  x1, y1, x2, y2 = box.tolist()
 
490
  if x2 <= x1 or y2 <= y1:
491
  continue
 
492
  results.append(
493
  BoundingBox(
494
+ x1=max(0, min(orig_w, int(math.floor(x1)))),
495
+ y1=max(0, min(orig_h, int(math.floor(y1)))),
496
+ x2=max(0, min(orig_w, int(math.ceil(x2)))),
497
+ y2=max(0, min(orig_h, int(math.ceil(y2)))),
498
  cls_id=int(cls_id),
499
+ conf=float(max(0.0, min(1.0, conf))),
500
  )
501
  )
 
502
  return results
503
 
504
+ # ─── Single-view inference ────────────────────────────────────
505
+
506
+ def _predict_single(self, image_bgr: np.ndarray,
507
+ allow_pad: bool = True) -> list[BoundingBox]:
508
+ if image_bgr is None or not isinstance(image_bgr, np.ndarray):
509
+ raise ValueError("Invalid image input")
510
+ if image_bgr.dtype != np.uint8:
511
+ image_bgr = image_bgr.astype(np.uint8)
512
+
513
+ inp, meta = self._preprocess(image_bgr, allow_pad=allow_pad)
514
+ outputs = self.session.run(None, {self.input_name: inp})
515
+
516
+ ratio = float(meta["ratio"])
517
+ pad = (float(meta["pad_w"]), float(meta["pad_h"]))
518
+ orig_size = (int(meta["orig_w"]), int(meta["orig_h"]))
519
+ extra = (int(meta["extra_left"]), int(meta["extra_right"]))
520
+
521
+ return self._decode_v26_output(outputs[0], ratio, pad, orig_size, extra)
522
+
523
+ # ─── TTA inference ────────────────────────────────────────────
524
+
525
+ def _infer_single(self, image_bgr: ndarray) -> list[BoundingBox]:
526
+ """3-view TTA: original (no pad) + flip (no pad) + original (L/R padded).
527
+ All three views return ORIGINAL-image coords, then pooled."""
528
+ orig_h, orig_w = image_bgr.shape[:2]
529
+
530
+ # View 1: original, NO left/right padding
531
+ boxes_orig = self._predict_single(image_bgr, allow_pad=True)
532
+
533
+ # View 2: horizontal flip, NO left/right padding
534
+ flipped = cv2.flip(image_bgr, 1)
535
+ boxes_flip = self._predict_single(flipped, allow_pad=True)
536
+ w = image_bgr.shape[1]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
537
  boxes_flip = [
538
+ BoundingBox(x1=w - b.x2, y1=b.y1, x2=w - b.x1, y2=b.y2,
539
+ cls_id=b.cls_id, conf=b.conf)
 
 
540
  for b in boxes_flip
541
  ]
542
+
543
+ # View 3: original, WITH left/right padding (only meaningful if square)
544
+ # boxes_pad = self._predict_single(image_bgr, allow_pad=True)
545
+ # NOZZLE_CLS = self.class_names.index('nozzle') # == 2
546
+ # boxes_pad = [b for b in boxes_pad if b.cls_id != NOZZLE_CLS]
547
+
548
+ # # View 4: flip, WITH left/right padding (only meaningful if square)
549
+ # flipped = cv2.flip(image_bgr, 1)
550
+ # boxes_pad_flip = self._predict_single(flipped, allow_pad=True)
551
+ # boxes_pad_flip = [
552
+ # BoundingBox(x1=w - b.x2, y1=b.y1, x2=w - b.x1, y2=b.y2,
553
+ # cls_id=b.cls_id, conf=b.conf)
554
+ # for b in boxes_pad_flip
555
+ # ]
556
+ # NOZZLE_CLS = self.class_names.index('nozzle') # == 2
557
+ # boxes_pad_flip = [b for b in boxes_pad_flip if b.cls_id != NOZZLE_CLS]
558
+
559
+ all_boxes = boxes_orig + boxes_flip# + boxes_pad# + boxes_pad_flip
560
+ if not all_boxes:
561
  return []
562
 
563
+ coords = np.array([[b.x1, b.y1, b.x2, b.y2] for b in all_boxes], dtype=np.float32)
 
 
564
  scores = np.array([b.conf for b in all_boxes], dtype=np.float32)
565
  cls_ids = np.array([b.cls_id for b in all_boxes], dtype=np.int32)
566
 
567
+ # Per-class NMS (uses per-class IoU thresholds)
568
+ hard_keep = self._per_class_hard_nms(coords, scores, cls_ids)
569
  if len(hard_keep) == 0:
570
  return []
571
+
572
  if len(hard_keep) > self.max_det:
573
  top = np.argsort(-scores[hard_keep])[: self.max_det]
574
  hard_keep = hard_keep[top]
575
+
576
+ # For confidence boost, use average IoU threshold (or could use median)
577
+ avg_iou = float(np.mean(self._iou_thres_array))
578
  boosted = self._max_score_per_cluster(
579
  coords[hard_keep], cls_ids[hard_keep],
580
+ coords, scores, cls_ids, avg_iou,
581
  )
582
 
583
  kept_coords = coords[hard_keep]
584
  kept_cls = cls_ids[hard_keep]
585
+
586
+ # Cross-class dedup
587
  if len(kept_coords) > 1:
588
  kept_coords, boosted, kept_cls = self._cross_class_dedup_op(
589
  kept_coords, boosted, kept_cls, self.cross_iou_thresh
590
  )
591
 
592
+ out_boxes = []
593
+ for j in range(len(kept_coords)):
594
+ x1, y1, x2, y2 = kept_coords[j].tolist()
595
+ if x2 <= x1 or y2 <= y1:
596
+ continue
597
+ out_boxes.append(
598
+ BoundingBox(
599
+ x1=max(0, min(orig_w, int(math.floor(x1)))),
600
+ y1=max(0, min(orig_h, int(math.floor(y1)))),
601
+ x2=max(0, min(orig_w, int(math.ceil(x2)))),
602
+ y2=max(0, min(orig_h, int(math.ceil(y2)))),
603
+ cls_id=int(kept_cls[j]),
604
+ conf=float(max(0.0, min(1.0, boosted[j]))),
605
+ )
606
  )
607
+ return out_boxes
 
608
 
609
+ # ─── Public API ───────────────────────────────────────────────
610
+
611
  def predict_batch(
612
  self,
613
  batch_images: list[ndarray],
 
615
  n_keypoints: int,
616
  ) -> list[TVFrameResult]:
617
  results: list[TVFrameResult] = []
618
+ for idx, image in enumerate(batch_images):
 
619
  try:
620
+ boxes = self._infer_single(image)
 
 
 
621
  except Exception as e:
622
+ print(f"Inference failed for frame {offset + idx}: {e}")
623
  boxes = []
624
+ keypoints = [(0, 0) for _ in range(max(0, int(n_keypoints)))]
625
  results.append(
626
  TVFrameResult(
627
+ frame_id=offset + idx,
628
  boxes=boxes,
629
+ keypoints=keypoints,
630
  )
631
  )
 
632
  return results
weights.onnx CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:1f9c13ee4403ddbbccd9d9707a151b1474cbb33124a2ef2965fa57d1e821ed03
3
- size 19287011
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:5eff40e23f79ec4d26d8639de704fdc432abab52f0fc44ec908037e9a2316824
3
+ size 20833918