SuperBitDev commited on
Commit
f3fe0f9
·
verified ·
1 Parent(s): 4952e2f

Upload folder using huggingface_hub

Browse files
Files changed (3) hide show
  1. chute_config.yml +8 -3
  2. miner.py +261 -219
  3. weights.onnx +2 -2
chute_config.yml CHANGED
@@ -7,9 +7,14 @@ Image:
7
 
8
  NodeSelector:
9
  gpu_count: 1
10
- min_vram_gb_per_gpu: 24
11
- min_memory_gb: 32
12
- min_cpu_count: 32
 
 
 
 
 
13
 
14
  Chute:
15
  timeout_seconds: 900
 
7
 
8
  NodeSelector:
9
  gpu_count: 1
10
+ min_vram_gb_per_gpu: 16
11
+ max_hourly_price_per_gpu: 0.5
12
+
13
+ exclude:
14
+ - "5090"
15
+ - b200
16
+ - h200
17
+ - mi300x
18
 
19
  Chute:
20
  timeout_seconds: 900
miner.py CHANGED
@@ -26,7 +26,7 @@ class TVFrameResult(BaseModel):
26
  class Miner:
27
  def __init__(self, path_hf_repo: Path) -> None:
28
  model_path = path_hf_repo / "weights.onnx"
29
- self.class_names = ["person"]
30
  print("ORT version:", ort.__version__)
31
 
32
  try:
@@ -67,31 +67,21 @@ class Miner:
67
  self.output_names = [output.name for output in self.session.get_outputs()]
68
  self.input_shape = self.session.get_inputs()[0].shape
69
 
 
70
  self.input_height = self._safe_dim(self.input_shape[2], default=1280)
71
  self.input_width = self._safe_dim(self.input_shape[3], default=1280)
72
 
73
- # ---------- Scoring-oriented thresholds ----------
74
- # Low threshold for candidate generation
75
- self.conf_thres = 0.25
76
-
77
- # High-confidence boxes can survive without TTA confirmation
78
- self.conf_high = 0.55
79
-
80
- # NMS threshold
81
- self.iou_thres = 0.50
82
-
83
- # TTA confirmation IoU
84
- self.tta_match_iou = 0.55
85
-
86
- self.max_det = 150
87
  self.use_tta = True
88
 
89
- # Box sanity filters
90
- self.min_box_area = 16 * 16
91
- self.min_w = 6
92
- self.min_h = 6
93
- self.max_aspect_ratio = 6.0
94
- self.max_box_area_ratio = 0.95
95
 
96
  print(f"✅ ONNX model loaded from: {model_path}")
97
  print(f"✅ ONNX providers: {self.session.get_providers()}")
@@ -113,6 +103,13 @@ class Miner:
113
  new_shape: tuple[int, int],
114
  color=(114, 114, 114),
115
  ) -> tuple[ndarray, float, tuple[float, float]]:
 
 
 
 
 
 
 
116
  h, w = image.shape[:2]
117
  new_w, new_h = new_shape
118
 
@@ -148,6 +145,14 @@ class Miner:
148
  def _preprocess(
149
  self, image: ndarray
150
  ) -> tuple[np.ndarray, float, tuple[float, float], tuple[int, int]]:
 
 
 
 
 
 
 
 
151
  orig_h, orig_w = image.shape[:2]
152
 
153
  img, ratio, pad = self._letterbox(
@@ -178,56 +183,93 @@ class Miner:
178
  out[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.0
179
  return out
180
 
181
- @staticmethod
182
- def _hard_nms(
183
  boxes: np.ndarray,
184
  scores: np.ndarray,
185
- iou_thresh: float,
186
- ) -> np.ndarray:
187
- if len(boxes) == 0:
188
- return np.array([], dtype=np.intp)
189
-
190
- boxes = np.asarray(boxes, dtype=np.float32)
191
- scores = np.asarray(scores, dtype=np.float32)
192
- order = np.argsort(scores)[::-1]
193
- keep = []
 
194
 
195
- while len(order) > 0:
196
- i = order[0]
197
- keep.append(i)
198
- if len(order) == 1:
199
- break
200
 
201
- rest = order[1:]
 
 
 
 
202
 
203
- xx1 = np.maximum(boxes[i, 0], boxes[rest, 0])
204
- yy1 = np.maximum(boxes[i, 1], boxes[rest, 1])
205
- xx2 = np.minimum(boxes[i, 2], boxes[rest, 2])
206
- yy2 = np.minimum(boxes[i, 3], boxes[rest, 3])
207
 
 
 
 
 
208
  inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
209
 
210
- area_i = np.maximum(0.0, (boxes[i, 2] - boxes[i, 0])) * np.maximum(0.0, (boxes[i, 3] - boxes[i, 1]))
211
- area_r = np.maximum(0.0, (boxes[rest, 2] - boxes[rest, 0])) * np.maximum(0.0, (boxes[rest, 3] - boxes[rest, 1]))
212
-
213
- iou = inter / (area_i + area_r - inter + 1e-7)
214
- order = rest[iou <= iou_thresh]
 
 
 
 
215
 
216
- return np.array(keep, dtype=np.intp)
 
217
 
218
  @staticmethod
219
- def _box_iou_one_to_many(box: np.ndarray, boxes: np.ndarray) -> np.ndarray:
220
- xx1 = np.maximum(box[0], boxes[:, 0])
221
- yy1 = np.maximum(box[1], boxes[:, 1])
222
- xx2 = np.minimum(box[2], boxes[:, 2])
223
- yy2 = np.minimum(box[3], boxes[:, 3])
224
-
225
- inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
226
-
227
- area_a = max(0.0, (box[2] - box[0]) * (box[3] - box[1]))
228
- area_b = np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) * np.maximum(0.0, boxes[:, 3] - boxes[:, 1])
229
-
230
- return inter / (area_a + area_b - inter + 1e-7)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
231
 
232
  def _filter_sane_boxes(
233
  self,
@@ -236,44 +278,69 @@ class Miner:
236
  cls_ids: np.ndarray,
237
  orig_size: tuple[int, int],
238
  ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
 
239
  if len(boxes) == 0:
240
  return boxes, scores, cls_ids
241
-
242
  orig_w, orig_h = orig_size
243
  image_area = float(orig_w * orig_h)
244
-
245
  keep = []
246
  for i, box in enumerate(boxes):
247
  x1, y1, x2, y2 = box.tolist()
248
  bw = x2 - x1
249
  bh = y2 - y1
250
-
251
  if bw <= 0 or bh <= 0:
252
  continue
253
- if bw < self.min_w or bh < self.min_h:
254
  continue
255
-
256
  area = bw * bh
257
  if area < self.min_box_area:
258
  continue
259
- if area > self.max_box_area_ratio * image_area:
260
  continue
261
-
262
  ar = max(bw / max(bh, 1e-6), bh / max(bw, 1e-6))
263
  if ar > self.max_aspect_ratio:
264
  continue
265
-
266
  keep.append(i)
267
-
268
  if not keep:
269
  return (
270
  np.empty((0, 4), dtype=np.float32),
271
  np.empty((0,), dtype=np.float32),
272
  np.empty((0,), dtype=np.int32),
273
  )
 
 
274
 
275
- keep = np.array(keep, dtype=np.intp)
276
- return boxes[keep], scores[keep], cls_ids[keep]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
277
 
278
  def _decode_final_dets(
279
  self,
@@ -281,7 +348,13 @@ class Miner:
281
  ratio: float,
282
  pad: tuple[float, float],
283
  orig_size: tuple[int, int],
 
284
  ) -> list[BoundingBox]:
 
 
 
 
 
285
  if preds.ndim == 3 and preds.shape[0] == 1:
286
  preds = preds[0]
287
 
@@ -292,13 +365,6 @@ class Miner:
292
  scores = preds[:, 4].astype(np.float32)
293
  cls_ids = preds[:, 5].astype(np.int32)
294
 
295
- # person only
296
- keep = cls_ids == 0
297
- boxes = boxes[keep]
298
- scores = scores[keep]
299
- cls_ids = cls_ids[keep]
300
-
301
- # candidate threshold
302
  keep = scores >= self.conf_thres
303
  boxes = boxes[keep]
304
  scores = scores[keep]
@@ -310,34 +376,51 @@ class Miner:
310
  pad_w, pad_h = pad
311
  orig_w, orig_h = orig_size
312
 
 
313
  boxes[:, [0, 2]] -= pad_w
314
  boxes[:, [1, 3]] -= pad_h
315
  boxes /= ratio
316
  boxes = self._clip_boxes(boxes, (orig_w, orig_h))
317
 
318
- boxes, scores, cls_ids = self._filter_sane_boxes(boxes, scores, cls_ids, orig_size)
 
 
 
319
  if len(boxes) == 0:
320
  return []
321
 
322
- keep_idx = self._hard_nms(boxes, scores, self.iou_thres)
323
- keep_idx = keep_idx[: self.max_det]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
324
 
325
- boxes = boxes[keep_idx]
326
- scores = scores[keep_idx]
327
- cls_ids = cls_ids[keep_idx]
328
 
329
- return [
330
- BoundingBox(
331
- x1=int(math.floor(box[0])),
332
- y1=int(math.floor(box[1])),
333
- x2=int(math.ceil(box[2])),
334
- y2=int(math.ceil(box[3])),
335
- cls_id=int(cls_id),
336
- conf=float(conf),
 
337
  )
338
- for box, conf, cls_id in zip(boxes, scores, cls_ids)
339
- if box[2] > box[0] and box[3] > box[1]
340
- ]
341
 
342
  def _decode_raw_yolo(
343
  self,
@@ -346,8 +429,15 @@ class Miner:
346
  pad: tuple[float, float],
347
  orig_size: tuple[int, int],
348
  ) -> list[BoundingBox]:
 
 
 
 
 
 
349
  if preds.ndim != 3:
350
  raise ValueError(f"Unexpected raw ONNX output shape: {preds.shape}")
 
351
  if preds.shape[0] != 1:
352
  raise ValueError(f"Unexpected batch dimension in raw output: {preds.shape}")
353
 
@@ -361,31 +451,14 @@ class Miner:
361
  raise ValueError(f"Unexpected normalized raw output shape: {preds.shape}")
362
 
363
  boxes_xywh = preds[:, :4].astype(np.float32)
364
- tail = preds[:, 4:].astype(np.float32)
365
-
366
- # Supports:
367
- # [x,y,w,h,score] single-class
368
- # [x,y,w,h,obj,cls] YOLO standard single-class
369
- # [x,y,w,h,obj,cls1,cls2,...] multi-class
370
- if tail.shape[1] == 1:
371
- scores = tail[:, 0]
372
- cls_ids = np.zeros(len(scores), dtype=np.int32)
373
- elif tail.shape[1] == 2:
374
- obj = tail[:, 0]
375
- cls_prob = tail[:, 1]
376
- scores = obj * cls_prob
377
  cls_ids = np.zeros(len(scores), dtype=np.int32)
378
  else:
379
- obj = tail[:, 0]
380
- class_probs = tail[:, 1:]
381
- cls_ids = np.argmax(class_probs, axis=1).astype(np.int32)
382
- cls_scores = class_probs[np.arange(len(class_probs)), cls_ids]
383
- scores = obj * cls_scores
384
-
385
- keep = cls_ids == 0
386
- boxes_xywh = boxes_xywh[keep]
387
- scores = scores[keep]
388
- cls_ids = cls_ids[keep]
389
 
390
  keep = scores >= self.conf_thres
391
  boxes_xywh = boxes_xywh[keep]
@@ -397,6 +470,12 @@ class Miner:
397
 
398
  boxes = self._xywh_to_xyxy(boxes_xywh)
399
 
 
 
 
 
 
 
400
  pad_w, pad_h = pad
401
  orig_w, orig_h = orig_size
402
 
@@ -405,29 +484,31 @@ class Miner:
405
  boxes /= ratio
406
  boxes = self._clip_boxes(boxes, (orig_w, orig_h))
407
 
408
- boxes, scores, cls_ids = self._filter_sane_boxes(boxes, scores, cls_ids, orig_size)
 
 
409
  if len(boxes) == 0:
410
  return []
411
 
412
- keep_idx = self._hard_nms(boxes, scores, self.iou_thres)
413
- keep_idx = keep_idx[: self.max_det]
 
414
 
415
- boxes = boxes[keep_idx]
416
- scores = scores[keep_idx]
417
- cls_ids = cls_ids[keep_idx]
418
 
419
- return [
420
- BoundingBox(
421
- x1=int(math.floor(box[0])),
422
- y1=int(math.floor(box[1])),
423
- x2=int(math.ceil(box[2])),
424
- y2=int(math.ceil(box[3])),
425
- cls_id=int(cls_id),
426
- conf=float(conf),
 
427
  )
428
- for box, conf, cls_id in zip(boxes, scores, cls_ids)
429
- if box[2] > box[0] and box[3] > box[1]
430
- ]
431
 
432
  def _postprocess(
433
  self,
@@ -436,12 +517,19 @@ class Miner:
436
  pad: tuple[float, float],
437
  orig_size: tuple[int, int],
438
  ) -> list[BoundingBox]:
 
 
 
 
 
439
  if output.ndim == 2 and output.shape[1] >= 6:
440
  return self._decode_final_dets(output, ratio, pad, orig_size)
441
 
442
- if output.ndim == 3 and output.shape[0] == 1 and output.shape[2] >= 6:
 
443
  return self._decode_final_dets(output, ratio, pad, orig_size)
444
 
 
445
  return self._decode_raw_yolo(output, ratio, pad, orig_size)
446
 
447
  def _predict_single(self, image: np.ndarray) -> list[BoundingBox]:
@@ -471,104 +559,58 @@ class Miner:
471
  det_output = outputs[0]
472
  return self._postprocess(det_output, ratio, pad, orig_size)
473
 
474
- def _merge_tta_consensus(
475
- self,
476
- boxes_orig: list[BoundingBox],
477
- boxes_flip: list[BoundingBox],
478
- ) -> list[BoundingBox]:
479
  """
480
- Keep:
481
- - any box with conf >= conf_high
482
- - low/medium-conf boxes only if confirmed across TTA views
483
- Then run final hard NMS.
484
  """
485
- if not boxes_orig and not boxes_flip:
486
- return []
487
-
488
- coords_o = np.array([[b.x1, b.y1, b.x2, b.y2] for b in boxes_orig], dtype=np.float32) if boxes_orig else np.empty((0, 4), dtype=np.float32)
489
- scores_o = np.array([b.conf for b in boxes_orig], dtype=np.float32) if boxes_orig else np.empty((0,), dtype=np.float32)
490
-
491
- coords_f = np.array([[b.x1, b.y1, b.x2, b.y2] for b in boxes_flip], dtype=np.float32) if boxes_flip else np.empty((0, 4), dtype=np.float32)
492
- scores_f = np.array([b.conf for b in boxes_flip], dtype=np.float32) if boxes_flip else np.empty((0,), dtype=np.float32)
493
-
494
- accepted_boxes = []
495
- accepted_scores = []
496
-
497
- # Original view candidates
498
- for i in range(len(coords_o)):
499
- score = scores_o[i]
500
- if score >= self.conf_high:
501
- accepted_boxes.append(coords_o[i])
502
- accepted_scores.append(score)
503
- elif len(coords_f) > 0:
504
- ious = self._box_iou_one_to_many(coords_o[i], coords_f)
505
- j = int(np.argmax(ious))
506
- if ious[j] >= self.tta_match_iou:
507
- fused_score = max(score, scores_f[j])
508
- accepted_boxes.append(coords_o[i])
509
- accepted_scores.append(fused_score)
510
-
511
- # Flipped-view high-confidence boxes that original missed
512
- for i in range(len(coords_f)):
513
- score = scores_f[i]
514
- if score < self.conf_high:
515
- continue
516
 
517
- if len(coords_o) == 0:
518
- accepted_boxes.append(coords_f[i])
519
- accepted_scores.append(score)
520
- continue
521
 
522
- ious = self._box_iou_one_to_many(coords_f[i], coords_o)
523
- if np.max(ious) < self.tta_match_iou:
524
- accepted_boxes.append(coords_f[i])
525
- accepted_scores.append(score)
 
 
 
 
526
 
527
- if not accepted_boxes:
 
528
  return []
529
 
530
- boxes = np.array(accepted_boxes, dtype=np.float32)
531
- scores = np.array(accepted_scores, dtype=np.float32)
532
-
533
- keep = self._hard_nms(boxes, scores, self.iou_thres)
534
- keep = keep[: self.max_det]
535
 
536
- out = []
537
- for idx in keep:
538
- x1, y1, x2, y2 = boxes[idx].tolist()
539
- out.append(
540
- BoundingBox(
541
- x1=int(math.floor(x1)),
542
- y1=int(math.floor(y1)),
543
- x2=int(math.ceil(x2)),
544
- y2=int(math.ceil(y2)),
545
- cls_id=0,
546
- conf=float(scores[idx]),
547
- )
548
- )
549
- return out
550
 
551
- def _predict_tta(self, image: np.ndarray) -> list[BoundingBox]:
552
- boxes_orig = self._predict_single(image)
553
 
554
- flipped = cv2.flip(image, 1)
555
- boxes_flip_raw = self._predict_single(flipped)
 
 
556
 
557
- w = image.shape[1]
558
- boxes_flip = [
559
  BoundingBox(
560
- x1=w - b.x2,
561
- y1=b.y1,
562
- x2=w - b.x1,
563
- y2=b.y2,
564
- cls_id=b.cls_id,
565
- conf=b.conf,
566
  )
567
- for b in boxes_flip_raw
568
  ]
569
 
570
- return self._merge_tta_consensus(boxes_orig, boxes_flip)
571
-
572
  def predict_batch(
573
  self,
574
  batch_images: list[ndarray],
 
26
  class Miner:
27
  def __init__(self, path_hf_repo: Path) -> None:
28
  model_path = path_hf_repo / "weights.onnx"
29
+ self.class_names = ['person']
30
  print("ORT version:", ort.__version__)
31
 
32
  try:
 
67
  self.output_names = [output.name for output in self.session.get_outputs()]
68
  self.input_shape = self.session.get_inputs()[0].shape
69
 
70
+ # Your export is fixed-size 1280, but we still read actual ONNX input shape first.
71
  self.input_height = self._safe_dim(self.input_shape[2], default=1280)
72
  self.input_width = self._safe_dim(self.input_shape[3], default=1280)
73
 
74
+ # Tuned for validator scoring: reduce FP (FALSE_POSITIVE pillar),
75
+ # preserve recall (MAP50, RECALL), improve precision.
76
+ self.conf_thres = 0.36 # Higher = fewer FP, slightly lower recall
77
+ self.iou_thres = 0.5 # Lower = suppress duplicate detections (FP)
78
+ self.max_det = 200 # Cap detections; sports ~20-30 persons
 
 
 
 
 
 
 
 
 
79
  self.use_tta = True
80
 
81
+ # Box sanity: filter tiny/spurious detections (common FP source)
82
+ self.min_box_area = 12 * 12 # ~144 px²
83
+ self.min_side = 8
84
+ self.max_aspect_ratio = 8.0
 
 
85
 
86
  print(f"✅ ONNX model loaded from: {model_path}")
87
  print(f"✅ ONNX providers: {self.session.get_providers()}")
 
103
  new_shape: tuple[int, int],
104
  color=(114, 114, 114),
105
  ) -> tuple[ndarray, float, tuple[float, float]]:
106
+ """
107
+ Resize with unchanged aspect ratio and pad to target shape.
108
+ Returns:
109
+ padded_image,
110
+ ratio,
111
+ (pad_w, pad_h) # half-padding
112
+ """
113
  h, w = image.shape[:2]
114
  new_w, new_h = new_shape
115
 
 
145
  def _preprocess(
146
  self, image: ndarray
147
  ) -> tuple[np.ndarray, float, tuple[float, float], tuple[int, int]]:
148
+ """
149
+ Preprocess for fixed-size ONNX export:
150
+ - enhance image quality (CLAHE, denoise, sharpen)
151
+ - letterbox to model input size
152
+ - BGR -> RGB
153
+ - normalize to [0,1]
154
+ - HWC -> NCHW float32
155
+ """
156
  orig_h, orig_w = image.shape[:2]
157
 
158
  img, ratio, pad = self._letterbox(
 
183
  out[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.0
184
  return out
185
 
186
+ def _soft_nms(
187
+ self,
188
  boxes: np.ndarray,
189
  scores: np.ndarray,
190
+ sigma: float = 0.5,
191
+ score_thresh: float = 0.01,
192
+ ) -> tuple[np.ndarray, np.ndarray]:
193
+ """
194
+ Soft-NMS: Gaussian decay of overlapping scores instead of hard removal.
195
+ Returns (kept_original_indices, updated_scores).
196
+ """
197
+ N = len(boxes)
198
+ if N == 0:
199
+ return np.array([], dtype=np.intp), np.array([], dtype=np.float32)
200
 
201
+ boxes = boxes.astype(np.float32, copy=True)
202
+ scores = scores.astype(np.float32, copy=True)
203
+ order = np.arange(N)
 
 
204
 
205
+ for i in range(N):
206
+ max_pos = i + int(np.argmax(scores[i:]))
207
+ boxes[[i, max_pos]] = boxes[[max_pos, i]]
208
+ scores[[i, max_pos]] = scores[[max_pos, i]]
209
+ order[[i, max_pos]] = order[[max_pos, i]]
210
 
211
+ if i + 1 >= N:
212
+ break
 
 
213
 
214
+ xx1 = np.maximum(boxes[i, 0], boxes[i + 1:, 0])
215
+ yy1 = np.maximum(boxes[i, 1], boxes[i + 1:, 1])
216
+ xx2 = np.minimum(boxes[i, 2], boxes[i + 1:, 2])
217
+ yy2 = np.minimum(boxes[i, 3], boxes[i + 1:, 3])
218
  inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
219
 
220
+ area_i = max(0.0, float(
221
+ (boxes[i, 2] - boxes[i, 0]) * (boxes[i, 3] - boxes[i, 1])
222
+ ))
223
+ areas_j = (
224
+ np.maximum(0.0, boxes[i + 1:, 2] - boxes[i + 1:, 0])
225
+ * np.maximum(0.0, boxes[i + 1:, 3] - boxes[i + 1:, 1])
226
+ )
227
+ iou = inter / (area_i + areas_j - inter + 1e-7)
228
+ scores[i + 1:] *= np.exp(-(iou ** 2) / sigma)
229
 
230
+ mask = scores > score_thresh
231
+ return order[mask], scores[mask]
232
 
233
  @staticmethod
234
+ def _hard_nms(
235
+ boxes: np.ndarray,
236
+ scores: np.ndarray,
237
+ iou_thresh: float,
238
+ ) -> np.ndarray:
239
+ """
240
+ Standard NMS: keep one box per overlapping cluster (the one with highest score).
241
+ Returns indices of kept boxes (into the boxes/scores arrays).
242
+ """
243
+ N = len(boxes)
244
+ if N == 0:
245
+ return np.array([], dtype=np.intp)
246
+ boxes = np.asarray(boxes, dtype=np.float32)
247
+ scores = np.asarray(scores, dtype=np.float32)
248
+ order = np.argsort(scores)[::-1]
249
+ keep: list[int] = []
250
+ suppressed = np.zeros(N, dtype=bool)
251
+ for i in range(N):
252
+ idx = order[i]
253
+ if suppressed[idx]:
254
+ continue
255
+ keep.append(idx)
256
+ bi = boxes[idx]
257
+ for k in range(i + 1, N):
258
+ jdx = order[k]
259
+ if suppressed[jdx]:
260
+ continue
261
+ bj = boxes[jdx]
262
+ xx1 = max(bi[0], bj[0])
263
+ yy1 = max(bi[1], bj[1])
264
+ xx2 = min(bi[2], bj[2])
265
+ yy2 = min(bi[3], bj[3])
266
+ inter = max(0.0, xx2 - xx1) * max(0.0, yy2 - yy1)
267
+ area_i = (bi[2] - bi[0]) * (bi[3] - bi[1])
268
+ area_j = (bj[2] - bj[0]) * (bj[3] - bj[1])
269
+ iou = inter / (area_i + area_j - inter + 1e-7)
270
+ if iou > iou_thresh:
271
+ suppressed[jdx] = True
272
+ return np.array(keep)
273
 
274
  def _filter_sane_boxes(
275
  self,
 
278
  cls_ids: np.ndarray,
279
  orig_size: tuple[int, int],
280
  ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
281
+ """Filter out tiny, degenerate, or implausible boxes (common FP)."""
282
  if len(boxes) == 0:
283
  return boxes, scores, cls_ids
 
284
  orig_w, orig_h = orig_size
285
  image_area = float(orig_w * orig_h)
 
286
  keep = []
287
  for i, box in enumerate(boxes):
288
  x1, y1, x2, y2 = box.tolist()
289
  bw = x2 - x1
290
  bh = y2 - y1
 
291
  if bw <= 0 or bh <= 0:
292
  continue
293
+ if bw < self.min_side or bh < self.min_side:
294
  continue
 
295
  area = bw * bh
296
  if area < self.min_box_area:
297
  continue
298
+ if area > 0.95 * image_area:
299
  continue
 
300
  ar = max(bw / max(bh, 1e-6), bh / max(bw, 1e-6))
301
  if ar > self.max_aspect_ratio:
302
  continue
 
303
  keep.append(i)
 
304
  if not keep:
305
  return (
306
  np.empty((0, 4), dtype=np.float32),
307
  np.empty((0,), dtype=np.float32),
308
  np.empty((0,), dtype=np.int32),
309
  )
310
+ k = np.array(keep, dtype=np.intp)
311
+ return boxes[k], scores[k], cls_ids[k]
312
 
313
+ @staticmethod
314
+ def _max_score_per_cluster(
315
+ coords: np.ndarray,
316
+ scores: np.ndarray,
317
+ keep_indices: np.ndarray,
318
+ iou_thresh: float,
319
+ ) -> np.ndarray:
320
+ """
321
+ For each kept box, return the max original score among itself and any
322
+ box that overlaps it with IOU >= iou_thresh (so TTA cluster keeps best conf).
323
+ """
324
+ n_keep = len(keep_indices)
325
+ if n_keep == 0:
326
+ return np.array([], dtype=np.float32)
327
+ out = np.empty(n_keep, dtype=np.float32)
328
+ coords = np.asarray(coords, dtype=np.float32)
329
+ scores = np.asarray(scores, dtype=np.float32)
330
+ for i in range(n_keep):
331
+ idx = keep_indices[i]
332
+ bi = coords[idx]
333
+ xx1 = np.maximum(bi[0], coords[:, 0])
334
+ yy1 = np.maximum(bi[1], coords[:, 1])
335
+ xx2 = np.minimum(bi[2], coords[:, 2])
336
+ yy2 = np.minimum(bi[3], coords[:, 3])
337
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
338
+ area_i = (bi[2] - bi[0]) * (bi[3] - bi[1])
339
+ areas_j = (coords[:, 2] - coords[:, 0]) * (coords[:, 3] - coords[:, 1])
340
+ iou = inter / (area_i + areas_j - inter + 1e-7)
341
+ in_cluster = iou >= iou_thresh
342
+ out[i] = float(np.max(scores[in_cluster]))
343
+ return out
344
 
345
  def _decode_final_dets(
346
  self,
 
348
  ratio: float,
349
  pad: tuple[float, float],
350
  orig_size: tuple[int, int],
351
+ apply_optional_dedup: bool = False,
352
  ) -> list[BoundingBox]:
353
+ """
354
+ Primary path:
355
+ expected output rows like [x1, y1, x2, y2, conf, cls_id]
356
+ in letterboxed input coordinates.
357
+ """
358
  if preds.ndim == 3 and preds.shape[0] == 1:
359
  preds = preds[0]
360
 
 
365
  scores = preds[:, 4].astype(np.float32)
366
  cls_ids = preds[:, 5].astype(np.int32)
367
 
 
 
 
 
 
 
 
368
  keep = scores >= self.conf_thres
369
  boxes = boxes[keep]
370
  scores = scores[keep]
 
376
  pad_w, pad_h = pad
377
  orig_w, orig_h = orig_size
378
 
379
+ # reverse letterbox
380
  boxes[:, [0, 2]] -= pad_w
381
  boxes[:, [1, 3]] -= pad_h
382
  boxes /= ratio
383
  boxes = self._clip_boxes(boxes, (orig_w, orig_h))
384
 
385
+ # Box sanity filter (reduces FP)
386
+ boxes, scores, cls_ids = self._filter_sane_boxes(
387
+ boxes, scores, cls_ids, orig_size
388
+ )
389
  if len(boxes) == 0:
390
  return []
391
 
392
+ # NMS to remove duplicates (model may output overlapping boxes)
393
+ if len(boxes) > 1:
394
+ if apply_optional_dedup:
395
+ keep_idx, scores = self._soft_nms(boxes, scores)
396
+ boxes = boxes[keep_idx]
397
+ cls_ids = cls_ids[keep_idx]
398
+ else:
399
+ keep_idx = self._hard_nms(boxes, scores, self.iou_thres)
400
+ keep_idx = keep_idx[: self.max_det]
401
+ boxes = boxes[keep_idx]
402
+ scores = scores[keep_idx]
403
+ cls_ids = cls_ids[keep_idx]
404
+
405
+ results: list[BoundingBox] = []
406
+ for box, conf, cls_id in zip(boxes, scores, cls_ids):
407
+ x1, y1, x2, y2 = box.tolist()
408
 
409
+ if x2 <= x1 or y2 <= y1:
410
+ continue
 
411
 
412
+ results.append(
413
+ BoundingBox(
414
+ x1=int(math.floor(x1)),
415
+ y1=int(math.floor(y1)),
416
+ x2=int(math.ceil(x2)),
417
+ y2=int(math.ceil(y2)),
418
+ cls_id=int(cls_id),
419
+ conf=float(conf),
420
+ )
421
  )
422
+
423
+ return results
 
424
 
425
  def _decode_raw_yolo(
426
  self,
 
429
  pad: tuple[float, float],
430
  orig_size: tuple[int, int],
431
  ) -> list[BoundingBox]:
432
+ """
433
+ Fallback path for raw YOLO predictions.
434
+ Supports common layouts:
435
+ - [1, C, N]
436
+ - [1, N, C]
437
+ """
438
  if preds.ndim != 3:
439
  raise ValueError(f"Unexpected raw ONNX output shape: {preds.shape}")
440
+
441
  if preds.shape[0] != 1:
442
  raise ValueError(f"Unexpected batch dimension in raw output: {preds.shape}")
443
 
 
451
  raise ValueError(f"Unexpected normalized raw output shape: {preds.shape}")
452
 
453
  boxes_xywh = preds[:, :4].astype(np.float32)
454
+ cls_part = preds[:, 4:].astype(np.float32)
455
+
456
+ if cls_part.shape[1] == 1:
457
+ scores = cls_part[:, 0]
 
 
 
 
 
 
 
 
 
458
  cls_ids = np.zeros(len(scores), dtype=np.int32)
459
  else:
460
+ cls_ids = np.argmax(cls_part, axis=1).astype(np.int32)
461
+ scores = cls_part[np.arange(len(cls_part)), cls_ids]
 
 
 
 
 
 
 
 
462
 
463
  keep = scores >= self.conf_thres
464
  boxes_xywh = boxes_xywh[keep]
 
470
 
471
  boxes = self._xywh_to_xyxy(boxes_xywh)
472
 
473
+ keep_idx = self._hard_nms(boxes, scores, self.iou_thres)
474
+ keep_idx = keep_idx[: self.max_det]
475
+ boxes = boxes[keep_idx]
476
+ scores = scores[keep_idx]
477
+ cls_ids = cls_ids[keep_idx]
478
+
479
  pad_w, pad_h = pad
480
  orig_w, orig_h = orig_size
481
 
 
484
  boxes /= ratio
485
  boxes = self._clip_boxes(boxes, (orig_w, orig_h))
486
 
487
+ boxes, scores, cls_ids = self._filter_sane_boxes(
488
+ boxes, scores, cls_ids, (orig_w, orig_h)
489
+ )
490
  if len(boxes) == 0:
491
  return []
492
 
493
+ results: list[BoundingBox] = []
494
+ for box, conf, cls_id in zip(boxes, scores, cls_ids):
495
+ x1, y1, x2, y2 = box.tolist()
496
 
497
+ if x2 <= x1 or y2 <= y1:
498
+ continue
 
499
 
500
+ results.append(
501
+ BoundingBox(
502
+ x1=int(math.floor(x1)),
503
+ y1=int(math.floor(y1)),
504
+ x2=int(math.ceil(x2)),
505
+ y2=int(math.ceil(y2)),
506
+ cls_id=int(cls_id),
507
+ conf=float(conf),
508
+ )
509
  )
510
+
511
+ return results
 
512
 
513
  def _postprocess(
514
  self,
 
517
  pad: tuple[float, float],
518
  orig_size: tuple[int, int],
519
  ) -> list[BoundingBox]:
520
+ """
521
+ Prefer final detections first.
522
+ Fallback to raw decode only if needed.
523
+ """
524
+ # final detections: [N,6]
525
  if output.ndim == 2 and output.shape[1] >= 6:
526
  return self._decode_final_dets(output, ratio, pad, orig_size)
527
 
528
+ # final detections: [1,N,6]
529
+ if output.ndim == 3 and output.shape[0] == 1 and output.shape[2] == 6:
530
  return self._decode_final_dets(output, ratio, pad, orig_size)
531
 
532
+ # fallback raw decode
533
  return self._decode_raw_yolo(output, ratio, pad, orig_size)
534
 
535
  def _predict_single(self, image: np.ndarray) -> list[BoundingBox]:
 
559
  det_output = outputs[0]
560
  return self._postprocess(det_output, ratio, pad, orig_size)
561
 
562
+ def _predict_tta(self, image: np.ndarray) -> list[BoundingBox]:
 
 
 
 
563
  """
564
+ Horizontal-flip TTA: merge original + flipped via hard NMS.
565
+ Boost confidence for consensus detections (both views agree) to improve
566
+ mAP: validator sorts by confidence, so higher conf for TP helps PR curve.
 
567
  """
568
+ boxes_orig = self._predict_single(image)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
569
 
570
+ flipped = cv2.flip(image, 1)
571
+ boxes_flip = self._predict_single(flipped)
 
 
572
 
573
+ w = image.shape[1]
574
+ boxes_flip = [
575
+ BoundingBox(
576
+ x1=w - b.x2, y1=b.y1, x2=w - b.x1, y2=b.y2,
577
+ cls_id=b.cls_id, conf=b.conf,
578
+ )
579
+ for b in boxes_flip
580
+ ]
581
 
582
+ all_boxes = boxes_orig + boxes_flip
583
+ if len(all_boxes) == 0:
584
  return []
585
 
586
+ coords = np.array(
587
+ [[b.x1, b.y1, b.x2, b.y2] for b in all_boxes], dtype=np.float32
588
+ )
589
+ scores = np.array([b.conf for b in all_boxes], dtype=np.float32)
 
590
 
591
+ hard_keep = self._hard_nms(coords, scores, self.iou_thres)
592
+ if len(hard_keep) == 0:
593
+ return []
 
 
 
 
 
 
 
 
 
 
 
594
 
595
+ hard_keep = hard_keep[: self.max_det]
 
596
 
597
+ # Boost confidence when both views agree (overlapping detections)
598
+ boosted = self._max_score_per_cluster(
599
+ coords, scores, hard_keep, self.iou_thres
600
+ )
601
 
602
+ return [
 
603
  BoundingBox(
604
+ x1=all_boxes[i].x1,
605
+ y1=all_boxes[i].y1,
606
+ x2=all_boxes[i].x2,
607
+ y2=all_boxes[i].y2,
608
+ cls_id=all_boxes[i].cls_id,
609
+ conf=float(boosted[j]),
610
  )
611
+ for j, i in enumerate(hard_keep)
612
  ]
613
 
 
 
614
  def predict_batch(
615
  self,
616
  batch_images: list[ndarray],
weights.onnx CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:4a7fac9f674924560d7a4788e2e1a4733b1055e4380f9e72e0720ed9fb3d155b
3
- size 19404973
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:1765e801b4092cc19c61e2cd7531e8be29517a8960eba329638d562be6a73a4d
3
+ size 19437023