SuperBitDev commited on
Commit
100117e
·
verified ·
1 Parent(s): fc37e5b

Upload folder using huggingface_hub

Browse files
Files changed (2) hide show
  1. miner.py +354 -651
  2. weights.onnx +2 -2
miner.py CHANGED
@@ -8,6 +8,35 @@ from numpy import ndarray
8
  from pydantic import BaseModel
9
 
10
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
11
  class BoundingBox(BaseModel):
12
  x1: int
13
  y1: int
@@ -24,32 +53,26 @@ 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
  sess_options.intra_op_num_threads = 2
54
  sess_options.inter_op_num_threads = 1
55
  sess_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
@@ -60,21 +83,18 @@ class Miner:
60
  sess_options=sess_options,
61
  providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
62
  )
63
- print("Created ORT session with preferred CUDA provider list")
64
  except Exception as e:
65
- print(f"⚠️ CUDA session creation failed, falling back to CPU: {e}")
66
  self.session = ort.InferenceSession(
67
  str(model_path),
68
  sess_options=sess_options,
69
  providers=["CPUExecutionProvider"],
70
  )
71
-
72
  print("ORT session providers:", self.session.get_providers())
73
 
74
- # Build cls_remap: for each model-emit index i,
75
- # cls_remap[i] = self.class_names.index(model_class_order[i])
76
- # The model-side order comes from the ONNX metadata when available,
77
- # else falls back to the static _model_class_order.
78
  model_class_order = self._read_model_class_order()
79
  if model_class_order is None:
80
  model_class_order = list(self._model_class_order)
@@ -87,81 +107,84 @@ class Miner:
87
 
88
  for inp in self.session.get_inputs():
89
  print("INPUT:", inp.name, inp.shape, inp.type)
90
-
91
  for out in self.session.get_outputs():
92
  print("OUTPUT:", out.name, out.shape, out.type)
93
 
94
  self.input_name = self.session.get_inputs()[0].name
95
- self.output_names = [output.name for output in self.session.get_outputs()]
96
  self.input_shape = self.session.get_inputs()[0].shape
97
 
98
- # Match the ONNX input dtype (this export is FP16 -> needs float16 input).
99
  input_type = self.session.get_inputs()[0].type
100
  self.np_dtype = np.float16 if "float16" in input_type else np.float32
101
- print(f"ONNX input dtype: {input_type} -> numpy {self.np_dtype}")
102
-
103
- # ONNX is fixed-size 1408x1408 (v1 export); read actual shape to be safe.
104
- self.input_height = self._safe_dim(self.input_shape[2], default=1280)
105
- self.input_width = self._safe_dim(self.input_shape[3], default=1280)
106
-
107
- # Tuned for validator scoring (pillars: 0.6*map50 + 0.4*false_positive).
108
- # All values below are the measured optimum of a full TTA-off grid sweep
109
- # on the validator-style val split (tune_miner.py, car-wash-49-val1024
110
- # val, ALL 2476 crops, against the synthetic-crop model; composite
111
- # 0.8716 -> 0.8731) -- re-run the sweep after any retrain.
112
- self.iou_thres = 0.5 # Per-class NMS IoU; lower = stricter dedup
113
- self.cross_iou_thresh = 0.9 # Cross-class dedup IoU (suppress same physical object firing multiple classes)
114
  self.max_det = 200
115
- # TTA = a 2nd (flipped) forward pass. Doubles latency; off for the
116
- # CPU latency gate. Re-enable only if the latency budget allows.
117
- self.use_tta = False
118
-
119
- # conf thresholds: broom=0.38 drainage gate=0.45 nozzle=0.30 track=0.60
120
- # Per-class confidence thresholds.
121
- # Indexed by class_names order: [broom, drainage gate, nozzle, track].
122
- # Values are the 2476-crop TTA-off sweep optimum on the synthetic-crop
123
- # model: lowering drainage (0.30->0.22) and nozzle (0.45->0.38) recovers
124
- # recall (map50 0.836->0.847) while FP actually drops (0.925->0.913).
125
- self._conf_thres_array = np.array(
126
- [0.30, 0.30, 0.40, 0.30], dtype=np.float32
127
- )
128
- # Per-class rescue bonus: when a class has ZERO boxes passing the
129
- # threshold in a frame, its top-1 candidate is admitted when its score
130
- # is at least (per-class threshold - per-class bonus).
131
- # DISABLED (all zeros): the 2476-crop sweep (rescue_bonus=False won)
132
- # confirmed rescue admits more false positives than true positives
133
- # under the validator's FP pillar.
134
- self._bonus_array = np.array(
135
- [0.05, 0.05, 0.1, 0.05], dtype=np.float32
136
- )
137
-
138
- # Box sanity filter kept loose: car-wash `nozzle` boxes are tiny
139
- # (GT median ~290 px², smallest ~32 px²). Fire's 14x14/min_side 8
140
- # would delete valid nozzles, so thresholds are dropped here.
141
- self.min_box_area = 4 * 4 # 16 px²
142
  self.min_side = 3
143
  self.max_aspect_ratio = 12.0
 
144
 
145
- print(f"✅ ONNX model loaded from: {model_path}")
146
- print(f"✅ ONNX providers: {self.session.get_providers()}")
147
- print(f"✅ ONNX input: name={self.input_name}, shape={self.input_shape}")
 
 
148
 
 
 
 
149
  self._warmup()
150
 
 
 
151
  def _warmup(self, iters: int = 3) -> None:
152
  try:
153
  dummy = np.zeros((720, 1280, 3), dtype=np.uint8)
154
  for _ in range(max(1, iters)):
155
  self.predict_batch(batch_images=[dummy], offset=0, n_keypoints=0)
156
- print(f"warmup: {iters} dummy predict_batch call(s) done")
157
  except Exception as e:
158
- print(f"⚠️ warmup skipped: {e}")
159
 
160
  def __repr__(self) -> str:
161
- return (
162
- f"ONNXRuntime(session={type(self.session).__name__}, "
163
- f"providers={self.session.get_providers()})"
164
- )
165
 
166
  @staticmethod
167
  def _safe_dim(value, default: int) -> int:
@@ -169,14 +192,7 @@ class Miner:
169
 
170
  @staticmethod
171
  def _resolve_model_path(repo: Path) -> Path:
172
- """Locate the ONNX model in the repo dir.
173
-
174
- Prefers weights.onnx (FP16/FP32 export), then weights_int8.onnx (the
175
- training script's INT8-quantized export -- works as-is: quantization
176
- preserves the Ultralytics metadata and QDQ models take regular fp32
177
- input), then any other .onnx file. INT8 is the fallback when the FP16
178
- export exceeds the 30 MB deployment limit (e.g. yolo26m).
179
- """
180
  for name in ("weights.onnx", "weights_int8.onnx"):
181
  p = repo / name
182
  if p.exists():
@@ -190,99 +206,72 @@ class Miner:
190
  return repo / "weights.onnx" # let session creation raise the error
191
 
192
  def _read_model_class_order(self) -> list[str] | None:
193
- """Read the model's class order from Ultralytics ONNX metadata.
194
-
195
- Returns the class names ordered by model-emit index, or None when
196
- metadata is missing/unparsable or doesn't match `class_names` as a
197
- set (in which case the static _model_class_order fallback is used).
198
- """
199
  try:
200
  import ast
201
-
202
  meta = self.session.get_modelmeta().custom_metadata_map
203
- names = ast.literal_eval(meta["names"]) # e.g. {0: 'broom', ...}
204
- if isinstance(names, dict):
205
- order = [str(names[i]) for i in sorted(names)]
206
- else:
207
- order = [str(n) for n in names]
208
  except Exception as e:
209
  print(f"cls order: could not read ONNX names metadata ({e})")
210
  return None
211
  if sorted(order) != sorted(self.class_names):
212
- print(
213
- f"cls order: ONNX names {order} do not match expected classes "
214
- f"{self.class_names}; ignoring metadata"
215
- )
216
  return None
217
  return order
218
 
219
- def _letterbox(
220
- self,
221
- image: ndarray,
222
- new_shape: tuple[int, int],
223
- color=(114, 114, 114),
224
- ) -> tuple[ndarray, float, tuple[float, float]]:
225
- """
226
- Resize with unchanged aspect ratio and pad to target shape.
227
- Returns:
228
- padded_image,
229
- ratio,
230
- (pad_w, pad_h) # half-padding
231
- """
232
  h, w = image.shape[:2]
233
  new_w, new_h = new_shape
234
-
235
  ratio = min(new_w / w, new_h / h)
236
- resized_w = int(round(w * ratio))
237
- resized_h = int(round(h * ratio))
238
-
239
- if (resized_w, resized_h) != (w, h):
240
  interp = cv2.INTER_CUBIC if ratio > 1.0 else cv2.INTER_LINEAR
241
- image = cv2.resize(image, (resized_w, resized_h), interpolation=interp)
242
-
243
- dw = new_w - resized_w
244
- dh = new_h - resized_h
245
- dw /= 2.0
246
- dh /= 2.0
247
-
248
- left = int(round(dw - 0.1))
249
- right = int(round(dw + 0.1))
250
- top = int(round(dh - 0.1))
251
- bottom = int(round(dh + 0.1))
252
-
253
- padded = cv2.copyMakeBorder(
254
- image,
255
- top,
256
- bottom,
257
- left,
258
- right,
259
- borderType=cv2.BORDER_CONSTANT,
260
- value=color,
261
- )
262
  return padded, ratio, (dw, dh)
263
 
264
- def _preprocess(
265
- self, image: ndarray
266
- ) -> tuple[np.ndarray, float, tuple[float, float], tuple[int, int]]:
267
- """
268
- Preprocess for fixed-size ONNX export:
269
- - enhance image quality (CLAHE, denoise, sharpen)
270
- - letterbox to model input size
271
- - BGR -> RGB
272
- - normalize to [0,1]
273
- - HWC -> NCHW float32
274
  """
275
  orig_h, orig_w = image.shape[:2]
276
-
277
- img, ratio, pad = self._letterbox(
278
- image, (self.input_width, self.input_height)
279
- )
 
 
 
 
 
 
 
280
  img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
281
- img = (img.astype(np.float32) / 255.0)
282
  img = np.transpose(img, (2, 0, 1))[None, ...]
283
  img = np.ascontiguousarray(img, dtype=self.np_dtype)
 
 
 
 
284
 
285
- return img, ratio, pad, (orig_w, orig_h)
286
 
287
  @staticmethod
288
  def _clip_boxes(boxes: np.ndarray, image_size: tuple[int, int]) -> np.ndarray:
@@ -302,222 +291,81 @@ class Miner:
302
  out[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.0
303
  return out
304
 
305
- def _soft_nms(
306
- self,
307
- boxes: np.ndarray,
308
- scores: np.ndarray,
309
- sigma: float = 0.5,
310
- score_thresh: float = 0.01,
311
- ) -> tuple[np.ndarray, np.ndarray]:
312
- """
313
- Soft-NMS: Gaussian decay of overlapping scores instead of hard removal.
314
- Returns (kept_original_indices, updated_scores).
315
- """
316
- N = len(boxes)
317
- if N == 0:
318
- return np.array([], dtype=np.intp), np.array([], dtype=np.float32)
319
-
320
- boxes = boxes.astype(np.float32, copy=True)
321
- scores = scores.astype(np.float32, copy=True)
322
- order = np.arange(N)
323
-
324
- for i in range(N):
325
- max_pos = i + int(np.argmax(scores[i:]))
326
- boxes[[i, max_pos]] = boxes[[max_pos, i]]
327
- scores[[i, max_pos]] = scores[[max_pos, i]]
328
- order[[i, max_pos]] = order[[max_pos, i]]
329
-
330
- if i + 1 >= N:
331
  break
332
-
333
- xx1 = np.maximum(boxes[i, 0], boxes[i + 1:, 0])
334
- yy1 = np.maximum(boxes[i, 1], boxes[i + 1:, 1])
335
- xx2 = np.minimum(boxes[i, 2], boxes[i + 1:, 2])
336
- yy2 = np.minimum(boxes[i, 3], boxes[i + 1:, 3])
337
  inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
 
 
 
338
 
339
- area_i = max(0.0, float(
340
- (boxes[i, 2] - boxes[i, 0]) * (boxes[i, 3] - boxes[i, 1])
341
- ))
342
- areas_j = (
343
- np.maximum(0.0, boxes[i + 1:, 2] - boxes[i + 1:, 0])
344
- * np.maximum(0.0, boxes[i + 1:, 3] - boxes[i + 1:, 1])
345
- )
346
- iou = inter / (area_i + areas_j - inter + 1e-7)
347
- scores[i + 1:] *= np.exp(-(iou ** 2) / sigma)
348
-
349
- mask = scores > score_thresh
350
- return order[mask], scores[mask]
351
-
352
- @staticmethod
353
- def _hard_nms(
354
- boxes: np.ndarray,
355
- scores: np.ndarray,
356
- iou_thresh: float,
357
- ) -> np.ndarray:
358
- """
359
- Standard NMS: keep one box per overlapping cluster (the one with highest score).
360
- Returns indices of kept boxes (into the boxes/scores arrays).
361
- """
362
- N = len(boxes)
363
- if N == 0:
364
- return np.array([], dtype=np.intp)
365
- boxes = np.asarray(boxes, dtype=np.float32)
366
- scores = np.asarray(scores, dtype=np.float32)
367
- order = np.argsort(scores)[::-1]
368
- keep: list[int] = []
369
- suppressed = np.zeros(N, dtype=bool)
370
- for i in range(N):
371
- idx = order[i]
372
- if suppressed[idx]:
373
- continue
374
- keep.append(idx)
375
- bi = boxes[idx]
376
- for k in range(i + 1, N):
377
- jdx = order[k]
378
- if suppressed[jdx]:
379
- continue
380
- bj = boxes[jdx]
381
- xx1 = max(bi[0], bj[0])
382
- yy1 = max(bi[1], bj[1])
383
- xx2 = min(bi[2], bj[2])
384
- yy2 = min(bi[3], bj[3])
385
- inter = max(0.0, xx2 - xx1) * max(0.0, yy2 - yy1)
386
- area_i = (bi[2] - bi[0]) * (bi[3] - bi[1])
387
- area_j = (bj[2] - bj[0]) * (bj[3] - bj[1])
388
- iou = inter / (area_i + area_j - inter + 1e-7)
389
- if iou > iou_thresh:
390
- suppressed[jdx] = True
391
- return np.array(keep)
392
-
393
- def _per_class_hard_nms(
394
- self,
395
- boxes: np.ndarray,
396
- scores: np.ndarray,
397
- cls_ids: np.ndarray,
398
- iou_thresh: float,
399
- ) -> np.ndarray:
400
- """Hard NMS applied independently per class."""
401
  if len(boxes) == 0:
402
  return np.array([], dtype=np.intp)
403
  all_keep: list[int] = []
404
  for c in np.unique(cls_ids):
405
- mask = cls_ids == c
406
- indices = np.where(mask)[0]
407
- keep = self._hard_nms(boxes[mask], scores[mask], iou_thresh)
408
- all_keep.extend(indices[keep].tolist())
409
  all_keep.sort()
410
  return np.array(all_keep, dtype=np.intp)
411
 
412
- def _per_class_soft_nms(
413
- self,
414
- boxes: np.ndarray,
415
- scores: np.ndarray,
416
- cls_ids: np.ndarray,
417
- sigma: float = 0.5,
418
- score_thresh: float = 0.01,
419
- ) -> tuple[np.ndarray, np.ndarray]:
420
- """Soft NMS applied independently per class."""
421
- if len(boxes) == 0:
422
- return np.array([], dtype=np.intp), np.array([], dtype=np.float32)
423
- all_keep: list[int] = []
424
- all_scores: list[float] = []
425
- for c in np.unique(cls_ids):
426
- mask = cls_ids == c
427
- indices = np.where(mask)[0]
428
- keep, updated = self._soft_nms(boxes[mask], scores[mask], sigma, score_thresh)
429
- for k, s in zip(keep, updated):
430
- all_keep.append(int(indices[k]))
431
- all_scores.append(float(s))
432
- if not all_keep:
433
- return np.array([], dtype=np.intp), np.array([], dtype=np.float32)
434
- return np.array(all_keep, dtype=np.intp), np.array(all_scores, dtype=np.float32)
435
-
436
- def _filter_sane_boxes(
437
- self,
438
- boxes: np.ndarray,
439
- scores: np.ndarray,
440
- cls_ids: np.ndarray,
441
- orig_size: tuple[int, int],
442
- ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
443
- """Filter out tiny, degenerate, or implausible boxes (common FP)."""
444
- if len(boxes) == 0:
445
- return boxes, scores, cls_ids
446
- orig_w, orig_h = orig_size
447
- image_area = float(orig_w * orig_h)
448
- keep = []
449
- for i, box in enumerate(boxes):
450
- x1, y1, x2, y2 = box.tolist()
451
- bw = x2 - x1
452
- bh = y2 - y1
453
- if bw <= 0 or bh <= 0:
454
- continue
455
- if bw < self.min_side or bh < self.min_side:
456
- continue
457
- area = bw * bh
458
- if area < self.min_box_area:
459
- continue
460
- if area > 0.95 * image_area:
461
- continue
462
- ar = max(bw / max(bh, 1e-6), bh / max(bw, 1e-6))
463
- if ar > self.max_aspect_ratio:
464
- continue
465
- keep.append(i)
466
- if not keep:
467
- return (
468
- np.empty((0, 4), dtype=np.float32),
469
- np.empty((0,), dtype=np.float32),
470
- np.empty((0,), dtype=np.int32),
471
- )
472
- k = np.array(keep, dtype=np.intp)
473
- return boxes[k], scores[k], cls_ids[k]
474
-
475
  @staticmethod
476
- def _max_score_per_cluster(
477
- post_boxes: np.ndarray,
478
- post_cls: np.ndarray,
479
- full_boxes: np.ndarray,
480
- full_scores: np.ndarray,
481
- full_cls: np.ndarray,
482
- iou_thresh: float,
483
- ) -> np.ndarray:
484
- """For each kept (post-NMS) box, return the max score over the FULL
485
- candidate set among SAME-CLASS boxes with IoU >= iou_thresh.
486
-
487
- The previous version omitted the same-class constraint, which let a
488
- confident broom raise the score of a coincident nozzle (or vice
489
- versa) under TTA. That's a silent FP booster and is fixed here.
490
- """
491
  n = len(post_boxes)
492
  if n == 0:
493
  return np.empty(0, dtype=np.float32)
494
- full_areas = (np.maximum(0.0, full_boxes[:, 2] - full_boxes[:, 0]) *
495
- np.maximum(0.0, full_boxes[:, 3] - full_boxes[:, 1]))
496
- out = np.empty(n, dtype=np.float32)
497
- for i in range(n):
498
- bi = post_boxes[i]
499
- xx1 = np.maximum(bi[0], full_boxes[:, 0])
500
- yy1 = np.maximum(bi[1], full_boxes[:, 1])
501
- xx2 = np.minimum(bi[2], full_boxes[:, 2])
502
- yy2 = np.minimum(bi[3], full_boxes[:, 3])
503
- inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
504
- a_i = max(0.0, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
505
- iou = inter / (a_i + full_areas - inter + 1e-7)
506
- cluster = (iou >= iou_thresh) & (full_cls == post_cls[i])
507
- out[i] = float(np.max(full_scores[cluster])) if np.any(cluster) else 0.0
508
- return out
509
-
510
- def _conf_filter_mask(
511
- self, scores: np.ndarray, cls_ids: np.ndarray
512
- ) -> np.ndarray:
513
- """Boolean keep-mask: score >= per-class threshold, with a per-class
514
- rescue -- if a class has zero boxes passing, admit its top-1 candidate
515
- when its score >= (per-class threshold - per-class bonus).
516
- """
517
  if len(scores) == 0:
518
  return np.zeros(0, dtype=bool)
519
- thr = self._conf_thres_array[cls_ids]
520
- keep = scores >= thr
521
  for c in np.unique(cls_ids):
522
  b = float(self._bonus_array[c])
523
  if b <= 0.0:
@@ -527,25 +375,36 @@ class Miner:
527
  continue
528
  idx = np.where(cm)[0]
529
  top = int(idx[int(np.argmax(scores[idx]))])
530
- if scores[top] >= self._conf_thres_array[c] - b:
531
  keep[top] = True
532
  return keep
533
 
534
- def _cross_class_dedup_op(
535
- self,
536
- boxes: np.ndarray,
537
- scores: np.ndarray,
538
- cls_ids: np.ndarray,
539
- iou_thresh: float,
540
- ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
541
- """Remove near-duplicate boxes across classes.
542
-
543
- Order candidates by (score - per_class_threshold) margin, then by area;
544
- keep the highest, suppress every other box with IoU > iou_thresh. For
545
- car-wash this kills the common failure where water spray makes the
546
- model fire both `nozzle` and `track` on the same patch, or where a
547
- broom handle overlaps a drainage-gate detection.
548
- """
 
 
 
 
 
 
 
 
 
 
 
549
  n = len(boxes)
550
  if n <= 1:
551
  return boxes, scores, cls_ids
@@ -576,141 +435,84 @@ class Miner:
576
  keep_idx = np.array(keep, dtype=np.intp)
577
  return boxes[keep_idx], scores[keep_idx], cls_ids[keep_idx]
578
 
579
- def _per_view_pipeline(
580
- self,
581
- boxes: np.ndarray,
582
- scores: np.ndarray,
583
- cls_ids: np.ndarray,
584
- ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
585
- """Per-view post-processing: per-class NMS -> cap -> cross-class dedup."""
 
 
 
586
  if len(boxes) > 1:
587
- keep = self._per_class_hard_nms(boxes, scores, cls_ids, self.iou_thres)
588
  boxes, scores, cls_ids = boxes[keep], scores[keep], cls_ids[keep]
589
  if len(scores) > self.max_det:
590
  top = np.argsort(-scores)[: self.max_det]
591
  boxes, scores, cls_ids = boxes[top], scores[top], cls_ids[top]
592
  if len(boxes) > 1:
593
  boxes, scores, cls_ids = self._cross_class_dedup_op(
594
- boxes, scores, cls_ids, self.cross_iou_thresh
595
- )
596
  return boxes, scores, cls_ids
597
 
598
- def _decode_final_dets(
599
- self,
600
- preds: np.ndarray,
601
- ratio: float,
602
- pad: tuple[float, float],
603
- orig_size: tuple[int, int],
604
- apply_optional_dedup: bool = False,
605
- ) -> list[BoundingBox]:
606
- """
607
- Primary path:
608
- expected output rows like [x1, y1, x2, y2, conf, cls_id]
609
- in letterboxed input coordinates.
610
- """
611
- if preds.ndim == 3 and preds.shape[0] == 1:
612
- preds = preds[0]
613
-
614
- if preds.ndim != 2 or preds.shape[1] < 6:
615
- raise ValueError(f"Unexpected ONNX final-det output shape: {preds.shape}")
616
-
617
- boxes = preds[:, :4].astype(np.float32)
618
- scores = preds[:, 4].astype(np.float32)
619
- cls_ids = preds[:, 5].astype(np.int32)
620
- cls_ids = self.cls_remap[cls_ids]
621
 
622
- # Per-class confidence filter with rescue (replaces scalar threshold)
623
- keep = self._conf_filter_mask(scores, cls_ids)
624
- boxes = boxes[keep]
625
- scores = scores[keep]
626
- cls_ids = cls_ids[keep]
 
 
 
627
 
628
- if len(boxes) == 0:
629
- return []
630
-
631
- pad_w, pad_h = pad
632
- orig_w, orig_h = orig_size
633
-
634
- # reverse letterbox
635
  boxes[:, [0, 2]] -= pad_w
636
  boxes[:, [1, 3]] -= pad_h
637
  boxes /= ratio
 
 
 
 
 
 
 
 
638
  boxes = self._clip_boxes(boxes, (orig_w, orig_h))
 
639
 
640
- # Box sanity filter (reduces FP)
641
- boxes, scores, cls_ids = self._filter_sane_boxes(
642
- boxes, scores, cls_ids, orig_size
643
- )
644
- if len(boxes) == 0:
645
- return []
646
-
647
- if apply_optional_dedup and len(boxes) > 1:
648
- # Soft-NMS path preserved as a tunable option; default below.
649
- keep_idx, scores = self._per_class_soft_nms(boxes, scores, cls_ids)
650
- boxes = boxes[keep_idx]
651
- cls_ids = cls_ids[keep_idx]
652
- if len(scores) > self.max_det:
653
- top = np.argsort(-scores)[: self.max_det]
654
- boxes, scores, cls_ids = boxes[top], scores[top], cls_ids[top]
655
- if len(boxes) > 1:
656
- boxes, scores, cls_ids = self._cross_class_dedup_op(
657
- boxes, scores, cls_ids, self.cross_iou_thresh
658
- )
659
- else:
660
- # Default: per-class hard NMS -> cap -> cross-class dedup
661
- boxes, scores, cls_ids = self._per_view_pipeline(boxes, scores, cls_ids)
662
-
663
- results: list[BoundingBox] = []
664
- for box, conf, cls_id in zip(boxes, scores, cls_ids):
665
- x1, y1, x2, y2 = box.tolist()
666
-
667
- if x2 <= x1 or y2 <= y1:
668
- continue
669
-
670
- results.append(
671
- BoundingBox(
672
- x1=int(math.floor(x1)),
673
- y1=int(math.floor(y1)),
674
- x2=int(math.ceil(x2)),
675
- y2=int(math.ceil(y2)),
676
- cls_id=int(cls_id),
677
- conf=float(conf),
678
- )
679
- )
680
-
681
- return results
682
-
683
- def _decode_raw_yolo(
684
- self,
685
- preds: np.ndarray,
686
- ratio: float,
687
- pad: tuple[float, float],
688
- orig_size: tuple[int, int],
689
- ) -> list[BoundingBox]:
690
- """
691
- Fallback path for raw YOLO predictions.
692
- Supports common layouts:
693
- - [1, C, N]
694
- - [1, N, C]
695
- """
696
- if preds.ndim != 3:
697
- raise ValueError(f"Unexpected raw ONNX output shape: {preds.shape}")
698
-
699
- if preds.shape[0] != 1:
700
- raise ValueError(f"Unexpected batch dimension in raw output: {preds.shape}")
701
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
702
  preds = preds[0]
703
-
704
- # Normalize to [N, C]
705
  if preds.shape[0] <= 16 and preds.shape[1] > preds.shape[0]:
706
  preds = preds.T
707
-
708
  if preds.ndim != 2 or preds.shape[1] < 5:
709
- raise ValueError(f"Unexpected normalized raw output shape: {preds.shape}")
710
-
711
  boxes_xywh = preds[:, :4].astype(np.float32)
712
  cls_part = preds[:, 4:].astype(np.float32)
713
-
714
  if cls_part.shape[1] == 1:
715
  scores = cls_part[:, 0]
716
  cls_ids = np.zeros(len(scores), dtype=np.int32)
@@ -718,196 +520,97 @@ class Miner:
718
  cls_ids = np.argmax(cls_part, axis=1).astype(np.int32)
719
  scores = cls_part[np.arange(len(cls_part)), cls_ids]
720
  cls_ids = self.cls_remap[cls_ids]
721
-
722
- # Per-class confidence filter with rescue (replaces scalar threshold)
723
- keep = self._conf_filter_mask(scores, cls_ids)
724
- boxes_xywh = boxes_xywh[keep]
725
- scores = scores[keep]
726
- cls_ids = cls_ids[keep]
727
  if len(boxes_xywh) == 0:
728
- return []
729
-
730
  boxes = self._xywh_to_xyxy(boxes_xywh)
 
731
 
732
- # Order matches fire001 / _decode_final_dets:
733
- # unscale -> clip -> sanity filter -> per-view pipeline (NMS, cap, cross-class dedup).
734
- pad_w, pad_h = pad
735
- orig_w, orig_h = orig_size
736
- boxes[:, [0, 2]] -= pad_w
737
- boxes[:, [1, 3]] -= pad_h
738
- boxes /= ratio
739
- boxes = self._clip_boxes(boxes, (orig_w, orig_h))
740
 
741
- boxes, scores, cls_ids = self._filter_sane_boxes(
742
- boxes, scores, cls_ids, (orig_w, orig_h)
743
- )
744
- if len(boxes) == 0:
745
- return []
746
 
747
- boxes, scores, cls_ids = self._per_view_pipeline(boxes, scores, cls_ids)
 
 
 
 
 
 
 
 
 
748
 
 
 
 
749
  results: list[BoundingBox] = []
 
750
  for box, conf, cls_id in zip(boxes, scores, cls_ids):
751
  x1, y1, x2, y2 = box.tolist()
752
-
753
  if x2 <= x1 or y2 <= y1:
754
  continue
755
-
756
- results.append(
757
- BoundingBox(
758
- x1=int(math.floor(x1)),
759
- y1=int(math.floor(y1)),
760
- x2=int(math.ceil(x2)),
761
- y2=int(math.ceil(y2)),
762
- cls_id=int(cls_id),
763
- conf=float(conf),
764
- )
765
- )
766
-
767
  return results
768
 
769
- def _postprocess(
770
- self,
771
- output: np.ndarray,
772
- ratio: float,
773
- pad: tuple[float, float],
774
- orig_size: tuple[int, int],
775
- ) -> list[BoundingBox]:
776
- """
777
- Prefer final detections first.
778
- Fallback to raw decode only if needed.
779
- """
780
- # final detections: [N,6]
781
- if output.ndim == 2 and output.shape[1] >= 6:
782
- return self._decode_final_dets(output, ratio, pad, orig_size)
783
-
784
- # final detections: [1,N,6]
785
- if output.ndim == 3 and output.shape[0] == 1 and output.shape[2] == 6:
786
- return self._decode_final_dets(output, ratio, pad, orig_size)
787
-
788
- # fallback raw decode
789
- return self._decode_raw_yolo(output, ratio, pad, orig_size)
790
-
791
- def _predict_single(self, image: np.ndarray) -> list[BoundingBox]:
792
- if image is None:
793
- raise ValueError("Input image is None")
794
- if not isinstance(image, np.ndarray):
795
- raise TypeError(f"Input is not numpy array: {type(image)}")
796
- if image.ndim != 3:
797
- raise ValueError(f"Expected HWC image, got shape={image.shape}")
798
- if image.shape[0] <= 0 or image.shape[1] <= 0:
799
- raise ValueError(f"Invalid image shape={image.shape}")
800
- if image.shape[2] != 3:
801
- raise ValueError(f"Expected 3 channels, got shape={image.shape}")
802
-
803
- if image.dtype != np.uint8:
804
- image = image.astype(np.uint8)
805
-
806
- input_tensor, ratio, pad, orig_size = self._preprocess(image)
807
-
808
- expected_shape = (1, 3, self.input_height, self.input_width)
809
- if input_tensor.shape != expected_shape:
810
- raise ValueError(
811
- f"Bad input tensor shape={input_tensor.shape}, expected={expected_shape}"
812
- )
813
-
814
- outputs = self.session.run(self.output_names, {self.input_name: input_tensor})
815
- det_output = outputs[0]
816
- return self._postprocess(det_output, ratio, pad, orig_size)
817
-
818
- def _predict_tta(self, image: np.ndarray) -> list[BoundingBox]:
819
- """Horizontal-flip TTA.
820
-
821
- Strategy (ported from fire001):
822
- 1. Predict on original and on flipped image.
823
- 2. Map flipped boxes back to original coordinates.
824
- 3. Per-class hard NMS on the union.
825
- 4. For each kept box, compute the max SAME-CLASS score across the
826
- FULL union -- a high-confidence flipped detection raises a
827
- borderline original one, but never one of a different class.
828
- 5. Cross-class dedup to suppress same-physical-object multi-class.
829
- """
830
- boxes_orig = self._predict_single(image)
831
-
832
- flipped = cv2.flip(image, 1)
833
- boxes_flip = self._predict_single(flipped)
834
-
835
  w = image.shape[1]
836
- boxes_flip = [
837
- BoundingBox(
838
- x1=w - b.x2, y1=b.y1, x2=w - b.x1, y2=b.y2,
839
- cls_id=b.cls_id, conf=b.conf,
840
- )
841
- for b in boxes_flip
842
- ]
843
-
844
- all_boxes = boxes_orig + boxes_flip
845
- if len(all_boxes) == 0:
846
  return []
847
-
848
- coords = np.array(
849
- [[b.x1, b.y1, b.x2, b.y2] for b in all_boxes], dtype=np.float32
850
- )
851
- scores = np.array([b.conf for b in all_boxes], dtype=np.float32)
852
- cls_ids = np.array([b.cls_id for b in all_boxes], dtype=np.int32)
853
-
854
- hard_keep = self._per_class_hard_nms(coords, scores, cls_ids, self.iou_thres)
855
- if len(hard_keep) == 0:
856
  return []
857
- if len(hard_keep) > self.max_det:
858
- top = np.argsort(-scores[hard_keep])[: self.max_det]
859
- hard_keep = hard_keep[top]
860
-
861
- # Class-aware cluster-max score boost (fixes the silent cross-class
862
- # leak in the previous _max_score_per_cluster).
863
  boosted = self._max_score_per_cluster(
864
- coords[hard_keep], cls_ids[hard_keep],
865
- coords, scores, cls_ids, self.iou_thres,
866
- )
867
-
868
- kept_coords = coords[hard_keep]
869
- kept_cls = cls_ids[hard_keep]
870
- if len(kept_coords) > 1:
871
- kept_coords, boosted, kept_cls = self._cross_class_dedup_op(
872
- kept_coords, boosted, kept_cls, self.cross_iou_thresh
873
- )
874
-
875
- return [
876
- BoundingBox(
877
- x1=int(math.floor(kept_coords[j, 0])),
878
- y1=int(math.floor(kept_coords[j, 1])),
879
- x2=int(math.ceil(kept_coords[j, 2])),
880
- y2=int(math.ceil(kept_coords[j, 3])),
881
- cls_id=int(kept_cls[j]),
882
- conf=float(boosted[j]),
883
- )
884
- for j in range(len(kept_coords))
885
- ]
886
-
887
- def predict_batch(
888
- self,
889
- batch_images: list[ndarray],
890
- offset: int,
891
- n_keypoints: int,
892
- ) -> list[TVFrameResult]:
893
  results: list[TVFrameResult] = []
894
-
895
- for frame_number_in_batch, image in enumerate(batch_images):
896
  try:
897
- if self.use_tta:
898
- boxes = self._predict_tta(image)
899
- else:
900
- boxes = self._predict_single(image)
901
  except Exception as e:
902
- print(f"⚠️ Inference failed for frame {offset + frame_number_in_batch}: {e}")
903
  boxes = []
904
-
905
- results.append(
906
- TVFrameResult(
907
- frame_id=offset + frame_number_in_batch,
908
- boxes=boxes,
909
- keypoints=[(0, 0) for _ in range(max(0, int(n_keypoints)))],
910
- )
911
- )
912
-
913
- return results
 
8
  from pydantic import BaseModel
9
 
10
 
11
+ # =============================================================================
12
+ # BEST merged car-wash miner.
13
+ #
14
+ # Base: carwash001/washvision01 (robust model loading + correct coordinate math)
15
+ # Merged-in ideas from ScoreVisionCarWash (the one genuinely different rival):
16
+ # * vectorized _hard_nms -> O(n^2) numpy, not a Python loop
17
+ # * vectorized _max_score_per_cluster -> single IoU matrix, not per-box loop
18
+ # * per-class NMS-IoU array -> _iou_thres_array (was one global)
19
+ # * per-class min-area array -> _min_box_area_array (was one global)
20
+ # * pre_nms_topk -> bounds NMS cost on crowded frames
21
+ # * single-view cluster boost -> confidence recovery on the DEPLOYED
22
+ # (non-TTA) path, not just under TTA
23
+ # * square-domain hook -> the validator scores 1024x1024
24
+ # squished frames (orig_w == orig_h);
25
+ # detect that and (a) use a val-tuned
26
+ # threshold set, (b) optionally widen
27
+ # to partly de-squish.
28
+ # Kept from the robust base (ScoreVision lacks all of these):
29
+ # * FP16 input auto-detect, cls_remap from ONNX metadata, INT8 + raw-YOLO
30
+ # decode fallbacks, ORT thread pinning, "no CLAHE" (train/test match).
31
+ #
32
+ # DEFAULTS reproduce carwash001's current DEPLOYED behavior (per-class IoU all
33
+ # 0.5, cluster boost off, square-pad off, thresholds = the tuned set) so this is
34
+ # a safe drop-in. The NEW levers are exposed but neutral until you re-tune them
35
+ # on the true-domain val (car-wash-55-styled/valid) with tune_miner.py. Do NOT
36
+ # assume the new knobs help before that sweep -- they are opportunities, measured.
37
+ # =============================================================================
38
+
39
+
40
  class BoundingBox(BaseModel):
41
  x1: int
42
  y1: int
 
53
 
54
 
55
  class Miner:
56
+ def __init__(self, path_hf_repo: Path) -> None:
 
 
57
  model_path = self._resolve_model_path(path_hf_repo)
58
+ # Canonical class order every downstream consumer (validator,
59
+ # BoundingBox.cls_id) sees: 0=broom, 1=drainage gate, 2=nozzle, 3=track.
 
60
  self.class_names = ["broom", "drainage gate", "nozzle", "track"]
61
+ # Fallback model-emit order (used only when ONNX metadata is missing).
 
 
 
 
62
  self._model_class_order = ["broom", "drainage gate", "nozzle", "track"]
63
  print("ORT version:", ort.__version__)
64
 
65
  try:
66
  ort.preload_dlls()
67
+ print("onnxruntime.preload_dlls() success")
68
  except Exception as e:
69
+ print(f"preload_dlls failed: {e}")
70
 
71
  print("ORT available providers BEFORE session:", ort.get_available_providers())
72
 
73
  sess_options = ort.SessionOptions()
74
  sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
75
+ # Pin threads for the CPU latency gate (ScoreVision leaves these default).
76
  sess_options.intra_op_num_threads = 2
77
  sess_options.inter_op_num_threads = 1
78
  sess_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
 
83
  sess_options=sess_options,
84
  providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
85
  )
86
+ print("Created ORT session with preferred CUDA provider list")
87
  except Exception as e:
88
+ print(f"CUDA session creation failed, falling back to CPU: {e}")
89
  self.session = ort.InferenceSession(
90
  str(model_path),
91
  sess_options=sess_options,
92
  providers=["CPUExecutionProvider"],
93
  )
 
94
  print("ORT session providers:", self.session.get_providers())
95
 
96
+ # cls_remap[i] = self.class_names.index(model_class_order[i]); order comes
97
+ # from ONNX metadata when present, else the static fallback.
 
 
98
  model_class_order = self._read_model_class_order()
99
  if model_class_order is None:
100
  model_class_order = list(self._model_class_order)
 
107
 
108
  for inp in self.session.get_inputs():
109
  print("INPUT:", inp.name, inp.shape, inp.type)
 
110
  for out in self.session.get_outputs():
111
  print("OUTPUT:", out.name, out.shape, out.type)
112
 
113
  self.input_name = self.session.get_inputs()[0].name
114
+ self.output_names = [o.name for o in self.session.get_outputs()]
115
  self.input_shape = self.session.get_inputs()[0].shape
116
 
117
+ # Match the ONNX input dtype (FP16 export needs float16 input).
118
  input_type = self.session.get_inputs()[0].type
119
  self.np_dtype = np.float16 if "float16" in input_type else np.float32
120
+ print(f"ONNX input dtype: {input_type} -> numpy {self.np_dtype}")
121
+
122
+ # The miner MUST run at the ONNX's baked input size. Dropping this to
123
+ # 640 (ScoreVision) is a real latency lever but is an EXPORT decision
124
+ # (re-export + re-eval; risks nozzle recall), not a miner-side change.
125
+ self.input_height = self._safe_dim(self.input_shape[2], default=704)
126
+ self.input_width = self._safe_dim(self.input_shape[3], default=704)
127
+
128
+ # ── Post-processing config (tune on car-wash-55-styled/valid) ─────────
129
+ # Per-class NMS IoU. Default all 0.5 == carwash001's single global value.
130
+ # ScoreVision uses [0.6,0.7,0.5,0.7]; sweep before adopting.
131
+ self._iou_thres_array = np.array([0.5, 0.5, 0.5, 0.5], dtype=np.float32)
132
+ self.cross_iou_thresh = 0.9 # cross-class dedup IoU (same physical object, 2 classes)
133
  self.max_det = 200
134
+ self.pre_nms_topk = 1000 # cap candidates before NMS (crowded-frame speed guard)
135
+ self.use_tta = False # 2nd flipped pass; doubles latency -> off for CPU gate
136
+
137
+ # Single-view cluster boost: raise each survivor's conf to its same-class
138
+ # IoU-cluster max. OFF by default (can raise FP under the FP pillar);
139
+ # ScoreVision runs it on. Sweep on the true-domain val before enabling.
140
+ self.use_cluster_boost = True
141
+
142
+ # Per-class confidence thresholds. `_conf_thres_array` is the SQUARE /
143
+ # validator-eval set (every scored frame is 1024x1024, so this is the
144
+ # one that matters); `_extra` is a fallback for non-square inputs
145
+ # (warmup / any 16:9 frame). Default: both = the tuned set.
146
+ # Current tuned optimum (2476-crop TTA-off sweep): 0.836->0.847 map50,
147
+ # FP 0.925->0.913. Re-run tune_miner on car-wash-55-styled/valid.
148
+ self._conf_thres_array = np.array([0.22, 0.22, 0.38, 0.28], dtype=np.float32)
149
+ self._extra_conf_thres_array = np.array([0.25, 0.25, 0.45, 0.30], dtype=np.float32)
150
+
151
+ # Per-class rescue bonus: if a class has ZERO boxes passing, admit its
152
+ # top-1 when score >= (threshold - bonus). DISABLED (all zeros): the
153
+ # sweep showed rescue admits more FP than TP under the FP pillar.
154
+ self._bonus_array = np.array([0.02, 0.02, 0.03, 0.03], dtype=np.float32)
155
+
156
+ # Per-class min box area (px^2). nozzle boxes are tiny (GT median ~290,
157
+ # min ~32 px^2) so its floor stays low. Default loose; ScoreVision uses
158
+ # [144,144,4,64]. Max-area cap is a fraction of the frame.
159
+ self._min_box_area_array = np.array([16.0, 16.0, 4.0, 16.0], dtype=np.float32)
 
160
  self.min_side = 3
161
  self.max_aspect_ratio = 12.0
162
+ self.max_area_frac = 0.95
163
 
164
+ # Square-domain de-squish: widen a square (validator) frame by this
165
+ # fraction before letterboxing, then drop pad-center boxes and un-pad.
166
+ # 0.0 = off (our model already scores 0.916 on squished val without it;
167
+ # enabling changes aspect -> validate first). ScoreVision uses 0.05.
168
+ self.square_pad_frac = 0.0
169
 
170
+ self._avg_iou = float(np.mean(self._iou_thres_array))
171
+ print(f"ONNX model loaded from: {model_path}")
172
+ print(f"ONNX input: name={self.input_name}, shape={self.input_shape}")
173
  self._warmup()
174
 
175
+ # ── Setup helpers ────────────────────────────────────────────────────────
176
+
177
  def _warmup(self, iters: int = 3) -> None:
178
  try:
179
  dummy = np.zeros((720, 1280, 3), dtype=np.uint8)
180
  for _ in range(max(1, iters)):
181
  self.predict_batch(batch_images=[dummy], offset=0, n_keypoints=0)
182
+ print(f"warmup: {iters} dummy predict_batch call(s) done")
183
  except Exception as e:
184
+ print(f"warmup skipped: {e}")
185
 
186
  def __repr__(self) -> str:
187
+ return f"CarWashMiner(classes={len(self.class_names)}, providers={self.session.get_providers()})"
 
 
 
188
 
189
  @staticmethod
190
  def _safe_dim(value, default: int) -> int:
 
192
 
193
  @staticmethod
194
  def _resolve_model_path(repo: Path) -> Path:
195
+ """Prefer weights.onnx, then weights_int8.onnx, then any .onnx."""
 
 
 
 
 
 
 
196
  for name in ("weights.onnx", "weights_int8.onnx"):
197
  p = repo / name
198
  if p.exists():
 
206
  return repo / "weights.onnx" # let session creation raise the error
207
 
208
  def _read_model_class_order(self) -> list[str] | None:
209
+ """Read class order from Ultralytics ONNX `names` metadata, or None."""
 
 
 
 
 
210
  try:
211
  import ast
 
212
  meta = self.session.get_modelmeta().custom_metadata_map
213
+ names = ast.literal_eval(meta["names"])
214
+ order = ([str(names[i]) for i in sorted(names)] if isinstance(names, dict)
215
+ else [str(n) for n in names])
 
 
216
  except Exception as e:
217
  print(f"cls order: could not read ONNX names metadata ({e})")
218
  return None
219
  if sorted(order) != sorted(self.class_names):
220
+ print(f"cls order: ONNX names {order} != expected {self.class_names}; ignoring")
 
 
 
221
  return None
222
  return order
223
 
224
+ # ── Preprocessing ────────────────────────────────────────────────────────
225
+
226
+ def _letterbox(self, image: ndarray, new_shape: tuple[int, int],
227
+ color=(114, 114, 114)) -> tuple[ndarray, float, tuple[float, float]]:
 
 
 
 
 
 
 
 
 
228
  h, w = image.shape[:2]
229
  new_w, new_h = new_shape
 
230
  ratio = min(new_w / w, new_h / h)
231
+ rw, rh = int(round(w * ratio)), int(round(h * ratio))
232
+ if (rw, rh) != (w, h):
 
 
233
  interp = cv2.INTER_CUBIC if ratio > 1.0 else cv2.INTER_LINEAR
234
+ image = cv2.resize(image, (rw, rh), interpolation=interp)
235
+ dw = (new_w - rw) / 2.0
236
+ dh = (new_h - rh) / 2.0
237
+ left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
238
+ top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
239
+ padded = cv2.copyMakeBorder(image, top, bottom, left, right,
240
+ cv2.BORDER_CONSTANT, value=color)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
241
  return padded, ratio, (dw, dh)
242
 
243
+ def _preprocess(self, image: ndarray) -> tuple[np.ndarray, dict]:
244
+ """Letterbox to the ONNX input size; NO CLAHE/denoise/sharpen (any
245
+ enhancement the model wasn't trained on is a train/test mismatch that
246
+ HURTS accuracy and also risks the CPU latency gate).
247
+
248
+ Square-domain de-squish: when the frame is square (the validator's
249
+ 1024x1024 squished eval image) and square_pad_frac > 0, widen it with
250
+ gray bars first -- objects come out horizontally compressed in those
251
+ frames, and mild widening partly restores their aspect.
 
252
  """
253
  orig_h, orig_w = image.shape[:2]
254
+ extra_left = extra_right = 0
255
+ if self.square_pad_frac > 0.0 and orig_w == orig_h:
256
+ target_w = int(orig_w * (1.0 + self.square_pad_frac))
257
+ if target_w > orig_w:
258
+ total = target_w - orig_w
259
+ extra_left = total // 2
260
+ extra_right = total - extra_left
261
+ image = cv2.copyMakeBorder(image, 0, 0, extra_left, extra_right,
262
+ cv2.BORDER_CONSTANT, value=(114, 114, 114))
263
+
264
+ img, ratio, pad = self._letterbox(image, (self.input_width, self.input_height))
265
  img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
266
+ img = img.astype(np.float32) / 255.0
267
  img = np.transpose(img, (2, 0, 1))[None, ...]
268
  img = np.ascontiguousarray(img, dtype=self.np_dtype)
269
+ return img, {
270
+ "ratio": ratio, "pad": pad, "orig_size": (orig_w, orig_h),
271
+ "extra_left": extra_left, "extra_right": extra_right,
272
+ }
273
 
274
+ # ── Vectorized box ops ───────────────────────────────────────────────────
275
 
276
  @staticmethod
277
  def _clip_boxes(boxes: np.ndarray, image_size: tuple[int, int]) -> np.ndarray:
 
291
  out[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.0
292
  return out
293
 
294
+ @staticmethod
295
+ def _hard_nms(boxes: np.ndarray, scores: np.ndarray, iou_thresh: float) -> np.ndarray:
296
+ """Vectorized greedy NMS (ScoreVision): areas precomputed once, each
297
+ step is a single numpy IoU vector -- no inner Python loop."""
298
+ n = len(boxes)
299
+ if n == 0:
300
+ return np.array([], dtype=np.intp)
301
+ x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
302
+ areas = np.maximum(0.0, x2 - x1) * np.maximum(0.0, y2 - y1)
303
+ order = np.argsort(-scores)
304
+ keep = []
305
+ while order.size > 0:
306
+ i = int(order[0])
307
+ keep.append(i)
308
+ if order.size == 1:
 
 
 
 
 
 
 
 
 
 
 
309
  break
310
+ rest = order[1:]
311
+ xx1 = np.maximum(x1[i], x1[rest])
312
+ yy1 = np.maximum(y1[i], y1[rest])
313
+ xx2 = np.minimum(x2[i], x2[rest])
314
+ yy2 = np.minimum(y2[i], y2[rest])
315
  inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
316
+ iou = inter / (areas[i] + areas[rest] - inter + 1e-7)
317
+ order = rest[iou <= iou_thresh]
318
+ return np.array(keep, dtype=np.intp)
319
 
320
+ def _per_class_hard_nms(self, boxes: np.ndarray, scores: np.ndarray,
321
+ cls_ids: np.ndarray) -> np.ndarray:
322
+ """Per-class NMS using the per-class IoU thresholds."""
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
323
  if len(boxes) == 0:
324
  return np.array([], dtype=np.intp)
325
  all_keep: list[int] = []
326
  for c in np.unique(cls_ids):
327
+ idx = np.where(cls_ids == c)[0]
328
+ cls_iou = float(self._iou_thres_array[c])
329
+ keep = self._hard_nms(boxes[idx], scores[idx], cls_iou)
330
+ all_keep.extend(idx[keep].tolist())
331
  all_keep.sort()
332
  return np.array(all_keep, dtype=np.intp)
333
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
334
  @staticmethod
335
+ def _max_score_per_cluster(post_boxes: np.ndarray, post_cls: np.ndarray,
336
+ full_boxes: np.ndarray, full_scores: np.ndarray,
337
+ full_cls: np.ndarray, iou_thresh: float) -> np.ndarray:
338
+ """Each survivor's confidence -> max score in its SAME-CLASS IoU cluster.
339
+ Vectorized single (n_post x n_full) IoU matrix (ScoreVision)."""
 
 
 
 
 
 
 
 
 
 
340
  n = len(post_boxes)
341
  if n == 0:
342
  return np.empty(0, dtype=np.float32)
343
+ m = len(full_boxes)
344
+ if m == 0:
345
+ return np.zeros(n, dtype=np.float32)
346
+ pa = (np.maximum(0.0, post_boxes[:, 2] - post_boxes[:, 0]) *
347
+ np.maximum(0.0, post_boxes[:, 3] - post_boxes[:, 1]))
348
+ fa = (np.maximum(0.0, full_boxes[:, 2] - full_boxes[:, 0]) *
349
+ np.maximum(0.0, full_boxes[:, 3] - full_boxes[:, 1]))
350
+ xx1 = np.maximum(post_boxes[:, 0][:, None], full_boxes[:, 0][None, :])
351
+ yy1 = np.maximum(post_boxes[:, 1][:, None], full_boxes[:, 1][None, :])
352
+ xx2 = np.minimum(post_boxes[:, 2][:, None], full_boxes[:, 2][None, :])
353
+ yy2 = np.minimum(post_boxes[:, 3][:, None], full_boxes[:, 3][None, :])
354
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
355
+ iou = inter / (pa[:, None] + fa[None, :] - inter + 1e-7)
356
+ mask = (iou >= iou_thresh) & (post_cls[:, None] == full_cls[None, :])
357
+ tiled = np.where(mask, full_scores[None, :], -np.inf)
358
+ out = tiled.max(axis=1)
359
+ out[~np.isfinite(out)] = 0.0
360
+ return out.astype(np.float32)
361
+
362
+ def _conf_filter_mask(self, scores: np.ndarray, cls_ids: np.ndarray,
363
+ extra_left: int) -> np.ndarray:
364
+ """Per-class threshold (square vs non-square set) + per-class rescue."""
 
365
  if len(scores) == 0:
366
  return np.zeros(0, dtype=bool)
367
+ thr_arr = self._extra_conf_thres_array if extra_left > 0 else self._conf_thres_array
368
+ keep = scores >= thr_arr[cls_ids]
369
  for c in np.unique(cls_ids):
370
  b = float(self._bonus_array[c])
371
  if b <= 0.0:
 
375
  continue
376
  idx = np.where(cm)[0]
377
  top = int(idx[int(np.argmax(scores[idx]))])
378
+ if scores[top] >= float(self._conf_thres_array[c]) - b:
379
  keep[top] = True
380
  return keep
381
 
382
+ def _filter_sane_boxes(self, boxes: np.ndarray, scores: np.ndarray,
383
+ cls_ids: np.ndarray, orig_size: tuple[int, int]
384
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
385
+ """Vectorized per-class min-area / max-area / min-side / aspect filter."""
386
+ if len(boxes) == 0:
387
+ return boxes, scores, cls_ids
388
+ orig_w, orig_h = orig_size
389
+ image_area = float(orig_w * orig_h)
390
+ bw = np.maximum(0.0, boxes[:, 2] - boxes[:, 0])
391
+ bh = np.maximum(0.0, boxes[:, 3] - boxes[:, 1])
392
+ area = bw * bh
393
+ ar = np.maximum(bw / np.maximum(bh, 1e-6), bh / np.maximum(bw, 1e-6))
394
+ keep = (
395
+ (bw >= self.min_side) & (bh >= self.min_side) &
396
+ (area >= self._min_box_area_array[cls_ids]) &
397
+ (area <= self.max_area_frac * image_area) &
398
+ (ar <= self.max_aspect_ratio)
399
+ )
400
+ return boxes[keep], scores[keep], cls_ids[keep]
401
+
402
+ def _cross_class_dedup_op(self, boxes: np.ndarray, scores: np.ndarray,
403
+ cls_ids: np.ndarray, iou_thresh: float
404
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
405
+ """Suppress near-duplicate boxes ACROSS classes (same physical object
406
+ firing 2 classes, e.g. spray -> nozzle+track). Order by conf-margin then
407
+ area; keep highest, drop IoU>thresh others."""
408
  n = len(boxes)
409
  if n <= 1:
410
  return boxes, scores, cls_ids
 
435
  keep_idx = np.array(keep, dtype=np.intp)
436
  return boxes[keep_idx], scores[keep_idx], cls_ids[keep_idx]
437
 
438
+ def _per_view_pipeline(self, boxes: np.ndarray, scores: np.ndarray,
439
+ cls_ids: np.ndarray, orig_size: tuple[int, int]
440
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
441
+ """sane filter -> top-k cap -> per-class NMS -> max_det cap -> cross-class dedup."""
442
+ boxes, scores, cls_ids = self._filter_sane_boxes(boxes, scores, cls_ids, orig_size)
443
+ if len(boxes) == 0:
444
+ return boxes, scores, cls_ids
445
+ if len(scores) > self.pre_nms_topk:
446
+ top = np.argpartition(-scores, self.pre_nms_topk)[: self.pre_nms_topk]
447
+ boxes, scores, cls_ids = boxes[top], scores[top], cls_ids[top]
448
  if len(boxes) > 1:
449
+ keep = self._per_class_hard_nms(boxes, scores, cls_ids)
450
  boxes, scores, cls_ids = boxes[keep], scores[keep], cls_ids[keep]
451
  if len(scores) > self.max_det:
452
  top = np.argsort(-scores)[: self.max_det]
453
  boxes, scores, cls_ids = boxes[top], scores[top], cls_ids[top]
454
  if len(boxes) > 1:
455
  boxes, scores, cls_ids = self._cross_class_dedup_op(
456
+ boxes, scores, cls_ids, self.cross_iou_thresh)
 
457
  return boxes, scores, cls_ids
458
 
459
+ # ── Coordinate un-mapping (shared by both decode paths) ──────────────────
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
460
 
461
+ def _unmap_and_finish(self, boxes: np.ndarray, scores: np.ndarray,
462
+ cls_ids: np.ndarray, meta: dict
463
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
464
+ """Undo letterbox + square-pad, drop pad-center boxes, clip, post-process."""
465
+ pad_w, pad_h = meta["pad"]
466
+ ratio = meta["ratio"]
467
+ orig_w, orig_h = meta["orig_size"]
468
+ extra_left, extra_right = meta["extra_left"], meta["extra_right"]
469
 
 
 
 
 
 
 
 
470
  boxes[:, [0, 2]] -= pad_w
471
  boxes[:, [1, 3]] -= pad_h
472
  boxes /= ratio
473
+ if extra_left:
474
+ boxes[:, [0, 2]] -= extra_left
475
+ if extra_left or extra_right:
476
+ cx = (boxes[:, 0] + boxes[:, 2]) * 0.5
477
+ inside = (cx >= 0) & (cx <= orig_w)
478
+ boxes, scores, cls_ids = boxes[inside], scores[inside], cls_ids[inside]
479
+ if len(boxes) == 0:
480
+ return boxes, scores, cls_ids
481
  boxes = self._clip_boxes(boxes, (orig_w, orig_h))
482
+ return self._per_view_pipeline(boxes, scores, cls_ids, (orig_w, orig_h))
483
 
484
+ # ── Decoding ─────────────────────────────────────────────────────────────
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
485
 
486
+ def _decode_final_dets(self, preds: np.ndarray, meta: dict
487
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
488
+ """End2end output rows [x1, y1, x2, y2, conf, cls_id] (NMS in graph)."""
489
+ empty = (np.empty((0, 4), np.float32), np.empty(0, np.float32), np.empty(0, np.int32))
490
+ if preds.ndim == 3 and preds.shape[0] == 1:
491
+ preds = preds[0]
492
+ if preds.ndim != 2 or preds.shape[1] < 6:
493
+ raise ValueError(f"Unexpected final-det output shape: {preds.shape}")
494
+ boxes = preds[:, :4].astype(np.float32)
495
+ scores = preds[:, 4].astype(np.float32)
496
+ cls_ids = self.cls_remap[preds[:, 5].astype(np.int32)]
497
+ keep = self._conf_filter_mask(scores, cls_ids, meta["extra_left"])
498
+ boxes, scores, cls_ids = boxes[keep], scores[keep], cls_ids[keep]
499
+ if len(boxes) == 0:
500
+ return empty
501
+ return self._unmap_and_finish(boxes, scores, cls_ids, meta)
502
+
503
+ def _decode_raw_yolo(self, preds: np.ndarray, meta: dict
504
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
505
+ """Raw YOLO output [1,C,N] or [1,N,C]; xywh + per-class scores."""
506
+ empty = (np.empty((0, 4), np.float32), np.empty(0, np.float32), np.empty(0, np.int32))
507
+ if preds.ndim != 3 or preds.shape[0] != 1:
508
+ raise ValueError(f"Unexpected raw output shape: {preds.shape}")
509
  preds = preds[0]
 
 
510
  if preds.shape[0] <= 16 and preds.shape[1] > preds.shape[0]:
511
  preds = preds.T
 
512
  if preds.ndim != 2 or preds.shape[1] < 5:
513
+ raise ValueError(f"Unexpected normalized raw shape: {preds.shape}")
 
514
  boxes_xywh = preds[:, :4].astype(np.float32)
515
  cls_part = preds[:, 4:].astype(np.float32)
 
516
  if cls_part.shape[1] == 1:
517
  scores = cls_part[:, 0]
518
  cls_ids = np.zeros(len(scores), dtype=np.int32)
 
520
  cls_ids = np.argmax(cls_part, axis=1).astype(np.int32)
521
  scores = cls_part[np.arange(len(cls_part)), cls_ids]
522
  cls_ids = self.cls_remap[cls_ids]
523
+ keep = self._conf_filter_mask(scores, cls_ids, meta["extra_left"])
524
+ boxes_xywh, scores, cls_ids = boxes_xywh[keep], scores[keep], cls_ids[keep]
 
 
 
 
525
  if len(boxes_xywh) == 0:
526
+ return empty
 
527
  boxes = self._xywh_to_xyxy(boxes_xywh)
528
+ return self._unmap_and_finish(boxes, scores, cls_ids, meta)
529
 
530
+ def _decode(self, output: np.ndarray, meta: dict
531
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
532
+ if output.ndim == 2 and output.shape[1] >= 6:
533
+ return self._decode_final_dets(output, meta)
534
+ if output.ndim == 3 and output.shape[0] == 1 and output.shape[2] == 6:
535
+ return self._decode_final_dets(output, meta)
536
+ return self._decode_raw_yolo(output, meta)
 
537
 
538
+ # ── Inference ────────────────────────────────────────────────────────────
 
 
 
 
539
 
540
+ def _predict_single_arrays(self, image: np.ndarray
541
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray, dict]:
542
+ if not isinstance(image, np.ndarray) or image.ndim != 3 or image.shape[2] != 3:
543
+ raise ValueError(f"Expected HWC BGR image, got {type(image)} {getattr(image,'shape',None)}")
544
+ if image.dtype != np.uint8:
545
+ image = image.astype(np.uint8)
546
+ inp, meta = self._preprocess(image)
547
+ outputs = self.session.run(self.output_names, {self.input_name: inp})
548
+ boxes, scores, cls_ids = self._decode(outputs[0], meta)
549
+ return boxes, scores, cls_ids, meta
550
 
551
+ @staticmethod
552
+ def _to_boxes(boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray,
553
+ orig_size: tuple[int, int]) -> list[BoundingBox]:
554
  results: list[BoundingBox] = []
555
+ orig_w, orig_h = orig_size
556
  for box, conf, cls_id in zip(boxes, scores, cls_ids):
557
  x1, y1, x2, y2 = box.tolist()
 
558
  if x2 <= x1 or y2 <= y1:
559
  continue
560
+ results.append(BoundingBox(
561
+ x1=max(0, min(orig_w, int(math.floor(x1)))),
562
+ y1=max(0, min(orig_h, int(math.floor(y1)))),
563
+ x2=max(0, min(orig_w, int(math.ceil(x2)))),
564
+ y2=max(0, min(orig_h, int(math.ceil(y2)))),
565
+ cls_id=int(cls_id),
566
+ conf=float(max(0.0, min(1.0, conf))),
567
+ ))
 
 
 
 
568
  return results
569
 
570
+ def _infer_single(self, image: np.ndarray) -> list[BoundingBox]:
571
+ boxes, scores, cls_ids, meta = self._predict_single_arrays(image)
572
+ if len(boxes) == 0:
573
+ return []
574
+ if self.use_cluster_boost and len(boxes) > 1:
575
+ scores = self._max_score_per_cluster(
576
+ boxes, cls_ids, boxes, scores, cls_ids, self._avg_iou)
577
+ return self._to_boxes(boxes, scores, cls_ids, meta["orig_size"])
578
+
579
+ def _infer_tta(self, image: np.ndarray) -> list[BoundingBox]:
580
+ """Horizontal-flip TTA: union of original + flipped, per-class NMS,
581
+ same-class cluster-max boost, cross-class dedup."""
582
+ b0, s0, c0, meta = self._predict_single_arrays(image)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
583
  w = image.shape[1]
584
+ bf, sf, cf, _ = self._predict_single_arrays(cv2.flip(image, 1))
585
+ if len(bf):
586
+ bf = bf.copy()
587
+ bf[:, [0, 2]] = w - bf[:, [2, 0]] # mirror x back to original coords
588
+ coords = np.concatenate([b0, bf], axis=0) if len(bf) else b0
589
+ scores = np.concatenate([s0, sf], axis=0) if len(sf) else s0
590
+ cls_ids = np.concatenate([c0, cf], axis=0) if len(cf) else c0
591
+ if len(coords) == 0:
 
 
592
  return []
593
+ keep = self._per_class_hard_nms(coords, scores, cls_ids)
594
+ if len(keep) == 0:
 
 
 
 
 
 
 
595
  return []
596
+ if len(keep) > self.max_det:
597
+ keep = keep[np.argsort(-scores[keep])[: self.max_det]]
 
 
 
 
598
  boosted = self._max_score_per_cluster(
599
+ coords[keep], cls_ids[keep], coords, scores, cls_ids, self._avg_iou)
600
+ kb, kc = coords[keep], cls_ids[keep]
601
+ if len(kb) > 1:
602
+ kb, boosted, kc = self._cross_class_dedup_op(kb, boosted, kc, self.cross_iou_thresh)
603
+ return self._to_boxes(kb, boosted, kc, meta["orig_size"])
604
+
605
+ def predict_batch(self, batch_images: list[ndarray], offset: int,
606
+ n_keypoints: int) -> list[TVFrameResult]:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
607
  results: list[TVFrameResult] = []
608
+ for idx, image in enumerate(batch_images):
 
609
  try:
610
+ boxes = self._infer_tta(image) if self.use_tta else self._infer_single(image)
 
 
 
611
  except Exception as e:
612
+ print(f"Inference failed for frame {offset + idx}: {e}")
613
  boxes = []
614
+ keypoints = [(0, 0) for _ in range(max(0, int(n_keypoints)))]
615
+ results.append(TVFrameResult(frame_id=offset + idx, boxes=boxes, keypoints=keypoints))
616
+ return results
 
 
 
 
 
 
 
weights.onnx CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:756f81aa467d482bb6508d172cc20c34fbb4cb7f15eeabb5538b858a1e561c94
3
- size 9824686
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:d5b41a6a181550440eda18f4af1f7a70e2e1edfa508fd5c29959e63ce824daeb
3
+ size 9842895