iotaminer commited on
Commit
d265b8a
·
verified ·
1 Parent(s): 23e58e4

scorevision: push artifact

Browse files
Files changed (1) hide show
  1. miner.py +235 -433
miner.py CHANGED
@@ -1,3 +1,13 @@
 
 
 
 
 
 
 
 
 
 
1
  from pathlib import Path
2
  import math
3
 
@@ -22,6 +32,7 @@ class TVFrameResult(BaseModel):
22
  boxes: list[BoundingBox]
23
  keypoints: list[tuple[int, int]]
24
 
 
25
  SIZE = 1280
26
 
27
 
@@ -37,17 +48,27 @@ class Miner:
37
  if ln.strip() and not ln.strip().startswith("#")
38
  ]
39
  else:
40
- self.class_names = ["person"]
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
 
@@ -57,9 +78,9 @@ class Miner:
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,
@@ -70,25 +91,35 @@ class Miner:
70
 
71
  for inp in self.session.get_inputs():
72
  print("INPUT:", inp.name, inp.shape, inp.type)
73
-
74
  for out in self.session.get_outputs():
75
  print("OUTPUT:", out.name, out.shape, out.type)
76
 
77
  self.input_name = self.session.get_inputs()[0].name
78
- self.output_names = [output.name for output in self.session.get_outputs()]
79
  self.input_shape = self.session.get_inputs()[0].shape
80
 
81
  self.input_height = self._safe_dim(self.input_shape[2], default=SIZE)
82
  self.input_width = self._safe_dim(self.input_shape[3], default=SIZE)
83
 
84
- self.conf_thres = 0.45
85
- self.iou_thres = 0.5
86
- self.max_det = 30
 
 
 
 
 
 
 
 
 
 
 
87
  self.use_tta = True
88
 
89
- print(f"ONNX model loaded from: {model_path}")
90
- print(f"ONNX providers: {self.session.get_providers()}")
91
- print(f"ONNX input: name={self.input_name}, shape={self.input_shape}")
92
 
93
  def __repr__(self) -> str:
94
  return (
@@ -100,73 +131,38 @@ class Miner:
100
  def _safe_dim(value, default: int) -> int:
101
  return value if isinstance(value, int) and value > 0 else default
102
 
 
103
  def _letterbox(
104
  self,
105
  image: ndarray,
106
  new_shape: tuple[int, int],
107
  color=(114, 114, 114),
108
  ) -> tuple[ndarray, float, tuple[float, float]]:
109
- """
110
- Resize with unchanged aspect ratio and pad to target shape.
111
- Returns:
112
- padded_image,
113
- ratio,
114
- (pad_w, pad_h) # half-padding
115
- """
116
  h, w = image.shape[:2]
117
  new_w, new_h = new_shape
118
-
119
  ratio = min(new_w / w, new_h / h)
120
  resized_w = int(round(w * ratio))
121
  resized_h = int(round(h * ratio))
122
-
123
  if (resized_w, resized_h) != (w, h):
124
  interp = cv2.INTER_CUBIC if ratio > 1.0 else cv2.INTER_LINEAR
125
  image = cv2.resize(image, (resized_w, resized_h), interpolation=interp)
126
-
127
- dw = new_w - resized_w
128
- dh = new_h - resized_h
129
- dw /= 2.0
130
- dh /= 2.0
131
-
132
  left = int(round(dw - 0.1))
133
  right = int(round(dw + 0.1))
134
  top = int(round(dh - 0.1))
135
  bottom = int(round(dh + 0.1))
136
-
137
  padded = cv2.copyMakeBorder(
138
- image,
139
- top,
140
- bottom,
141
- left,
142
- right,
143
- borderType=cv2.BORDER_CONSTANT,
144
- value=color,
145
  )
146
  return padded, ratio, (dw, dh)
147
 
148
- def _preprocess(
149
- self, image: ndarray
150
- ) -> tuple[np.ndarray, float, tuple[float, float], tuple[int, int]]:
151
- """
152
- Preprocess for fixed-size ONNX export:
153
- - enhance image quality (CLAHE, denoise, sharpen)
154
- - letterbox to model input size
155
- - BGR -> RGB
156
- - normalize to [0,1]
157
- - HWC -> NCHW float32
158
- """
159
- orig_h, orig_w = image.shape[:2]
160
-
161
- img, ratio, pad = self._letterbox(
162
- image, (self.input_width, self.input_height)
163
- )
164
- img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
165
- img = img.astype(np.float32) / 255.0
166
  img = np.transpose(img, (2, 0, 1))[None, ...]
167
- img = np.ascontiguousarray(img, dtype=np.float32)
168
-
169
- return img, ratio, pad, (orig_w, orig_h)
170
 
171
  @staticmethod
172
  def _clip_boxes(boxes: np.ndarray, image_size: tuple[int, int]) -> np.ndarray:
@@ -177,375 +173,244 @@ class Miner:
177
  boxes[:, 3] = np.clip(boxes[:, 3], 0, h - 1)
178
  return boxes
179
 
 
180
  @staticmethod
181
- def _xywh_to_xyxy(boxes: np.ndarray) -> np.ndarray:
182
- out = np.empty_like(boxes)
183
- out[:, 0] = boxes[:, 0] - boxes[:, 2] / 2.0
184
- out[:, 1] = boxes[:, 1] - boxes[:, 3] / 2.0
185
- out[:, 2] = boxes[:, 0] + boxes[:, 2] / 2.0
186
- out[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.0
187
- return out
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
188
 
189
  def _soft_nms(
190
  self,
191
  boxes: np.ndarray,
192
  scores: np.ndarray,
193
- sigma: float = 0.5,
194
  score_thresh: float = 0.01,
195
  ) -> tuple[np.ndarray, np.ndarray]:
196
- """
197
- Soft-NMS: Gaussian decay of overlapping scores instead of hard removal.
198
- Returns (kept_original_indices, updated_scores).
199
- """
200
  N = len(boxes)
201
  if N == 0:
202
  return np.array([], dtype=np.intp), np.array([], dtype=np.float32)
203
-
204
  boxes = boxes.astype(np.float32, copy=True)
205
  scores = scores.astype(np.float32, copy=True)
206
  order = np.arange(N)
207
-
208
  for i in range(N):
209
  max_pos = i + int(np.argmax(scores[i:]))
210
  boxes[[i, max_pos]] = boxes[[max_pos, i]]
211
  scores[[i, max_pos]] = scores[[max_pos, i]]
212
  order[[i, max_pos]] = order[[max_pos, i]]
213
-
214
  if i + 1 >= N:
215
  break
216
-
217
  xx1 = np.maximum(boxes[i, 0], boxes[i + 1:, 0])
218
  yy1 = np.maximum(boxes[i, 1], boxes[i + 1:, 1])
219
  xx2 = np.minimum(boxes[i, 2], boxes[i + 1:, 2])
220
  yy2 = np.minimum(boxes[i, 3], boxes[i + 1:, 3])
221
  inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
222
-
223
- area_i = max(0.0, float(
224
  (boxes[i, 2] - boxes[i, 0]) * (boxes[i, 3] - boxes[i, 1])
225
- ))
226
  areas_j = (
227
  np.maximum(0.0, boxes[i + 1:, 2] - boxes[i + 1:, 0])
228
  * np.maximum(0.0, boxes[i + 1:, 3] - boxes[i + 1:, 1])
229
  )
230
  iou = inter / (area_i + areas_j - inter + 1e-7)
231
  scores[i + 1:] *= np.exp(-(iou ** 2) / sigma)
232
-
233
  mask = scores > score_thresh
234
  return order[mask], scores[mask]
235
 
236
  @staticmethod
237
- def _hard_nms(
238
- boxes: np.ndarray,
239
- scores: np.ndarray,
240
- iou_thresh: float,
241
- ) -> np.ndarray:
242
- """
243
- Standard NMS: keep one box per overlapping cluster (the one with highest score).
244
- Returns indices of kept boxes (into the boxes/scores arrays).
245
- """
246
- N = len(boxes)
247
- if N == 0:
248
- return np.array([], dtype=np.intp)
249
- boxes = np.asarray(boxes, dtype=np.float32)
250
- scores = np.asarray(scores, dtype=np.float32)
251
- order = np.argsort(scores)[::-1]
252
- keep: list[int] = []
253
- suppressed = np.zeros(N, dtype=bool)
254
- for i in range(N):
255
- idx = order[i]
256
- if suppressed[idx]:
257
- continue
258
- keep.append(idx)
259
- bi = boxes[idx]
260
- for k in range(i + 1, N):
261
- jdx = order[k]
262
- if suppressed[jdx]:
263
- continue
264
- bj = boxes[jdx]
265
- xx1 = max(bi[0], bj[0])
266
- yy1 = max(bi[1], bj[1])
267
- xx2 = min(bi[2], bj[2])
268
- yy2 = min(bi[3], bj[3])
269
- inter = max(0.0, xx2 - xx1) * max(0.0, yy2 - yy1)
270
- area_i = (bi[2] - bi[0]) * (bi[3] - bi[1])
271
- area_j = (bj[2] - bj[0]) * (bj[3] - bj[1])
272
- iou = inter / (area_i + area_j - inter + 1e-7)
273
- if iou > iou_thresh:
274
- suppressed[jdx] = True
275
- return np.array(keep)
276
-
277
- @staticmethod
278
- def _max_score_per_cluster(
279
- coords: np.ndarray,
280
- scores: np.ndarray,
281
- keep_indices: np.ndarray,
282
- iou_thresh: float,
283
- ) -> np.ndarray:
284
- """
285
- For each kept box, return the max original score among itself and any
286
- box that overlaps it with IOU >= iou_thresh (so TTA cluster keeps best conf).
287
- """
288
- n_keep = len(keep_indices)
289
- if n_keep == 0:
290
- return np.array([], dtype=np.float32)
291
- out = np.empty(n_keep, dtype=np.float32)
292
- coords = np.asarray(coords, dtype=np.float32)
293
- scores = np.asarray(scores, dtype=np.float32)
294
- for i in range(n_keep):
295
- idx = keep_indices[i]
296
- bi = coords[idx]
297
- xx1 = np.maximum(bi[0], coords[:, 0])
298
- yy1 = np.maximum(bi[1], coords[:, 1])
299
- xx2 = np.minimum(bi[2], coords[:, 2])
300
- yy2 = np.minimum(bi[3], coords[:, 3])
301
- inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
302
- area_i = (bi[2] - bi[0]) * (bi[3] - bi[1])
303
- areas_j = (coords[:, 2] - coords[:, 0]) * (coords[:, 3] - coords[:, 1])
304
- iou = inter / (area_i + areas_j - inter + 1e-7)
305
- in_cluster = iou >= iou_thresh
306
- out[i] = float(np.max(scores[in_cluster]))
307
- return out
308
-
309
- def _decode_final_dets(
310
- self,
311
- preds: np.ndarray,
312
- ratio: float,
313
- pad: tuple[float, float],
314
- orig_size: tuple[int, int],
315
- apply_optional_dedup: bool = False,
316
- ) -> list[BoundingBox]:
317
- """
318
- Primary path:
319
- expected output rows like [x1, y1, x2, y2, conf, cls_id]
320
- in letterboxed input coordinates.
321
- """
322
- if preds.ndim == 3 and preds.shape[0] == 1:
323
- preds = preds[0]
324
-
325
- if preds.ndim != 2 or preds.shape[1] < 6:
326
- raise ValueError(f"Unexpected ONNX final-det output shape: {preds.shape}")
327
-
328
- boxes = preds[:, :4].astype(np.float32)
329
- scores = preds[:, 4].astype(np.float32)
330
- cls_ids = preds[:, 5].astype(np.int32)
331
-
332
- keep = scores >= self.conf_thres
333
- boxes = boxes[keep]
334
- scores = scores[keep]
335
- cls_ids = cls_ids[keep]
336
-
337
  if len(boxes) == 0:
338
- return []
339
-
340
- pad_w, pad_h = pad
341
- orig_w, orig_h = orig_size
342
-
343
- # reverse letterbox
344
- boxes[:, [0, 2]] -= pad_w
345
- boxes[:, [1, 3]] -= pad_h
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
346
  boxes /= ratio
347
- boxes = self._clip_boxes(boxes, (orig_w, orig_h))
348
-
349
- if apply_optional_dedup and len(boxes) > 1:
350
- keep_idx, scores = self._soft_nms(boxes, scores)
351
- boxes = boxes[keep_idx]
352
- cls_ids = cls_ids[keep_idx]
353
-
354
- results: list[BoundingBox] = []
355
- for box, conf, cls_id in zip(boxes, scores, cls_ids):
356
- x1, y1, x2, y2 = box.tolist()
357
 
358
- if x2 <= x1 or y2 <= y1:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
359
  continue
360
-
361
- results.append(
362
- BoundingBox(
363
- x1=int(math.floor(x1)),
364
- y1=int(math.floor(y1)),
365
- x2=int(math.ceil(x2)),
366
- y2=int(math.ceil(y2)),
367
- cls_id=int(cls_id),
368
- conf=float(conf),
369
- )
370
- )
371
-
372
- return results
373
-
374
- def _decode_raw_yolo(
375
- self,
376
- preds: np.ndarray,
377
- ratio: float,
378
- pad: tuple[float, float],
379
- orig_size: tuple[int, int],
380
- ) -> list[BoundingBox]:
381
- """
382
- Fallback path for raw YOLO predictions.
383
- Supports common layouts:
384
- - [1, C, N]
385
- - [1, N, C]
386
- """
387
- if preds.ndim != 3:
388
- raise ValueError(f"Unexpected raw ONNX output shape: {preds.shape}")
389
-
390
- if preds.shape[0] != 1:
391
- raise ValueError(f"Unexpected batch dimension in raw output: {preds.shape}")
392
-
393
- preds = preds[0]
394
-
395
- # Normalize to [N, C]
396
- if preds.shape[0] <= 16 and preds.shape[1] > preds.shape[0]:
397
- preds = preds.T
398
-
399
- if preds.ndim != 2 or preds.shape[1] < 5:
400
- raise ValueError(f"Unexpected normalized raw output shape: {preds.shape}")
401
-
402
- boxes_xywh = preds[:, :4].astype(np.float32)
403
- cls_part = preds[:, 4:].astype(np.float32)
404
-
405
- if cls_part.shape[1] == 1:
406
- scores = cls_part[:, 0]
407
- cls_ids = np.zeros(len(scores), dtype=np.int32)
408
- else:
409
- cls_ids = np.argmax(cls_part, axis=1).astype(np.int32)
410
- scores = cls_part[np.arange(len(cls_part)), cls_ids]
411
-
412
- keep = scores >= self.conf_thres
413
- boxes_xywh = boxes_xywh[keep]
414
- scores = scores[keep]
415
- cls_ids = cls_ids[keep]
416
-
417
- if len(boxes_xywh) == 0:
 
 
 
 
 
 
 
 
 
 
418
  return []
 
 
 
 
419
 
420
- boxes = self._xywh_to_xyxy(boxes_xywh)
421
- keep_idx, scores = self._soft_nms(boxes, scores)
422
- keep_idx = keep_idx[: self.max_det]
423
- scores = scores[: self.max_det]
424
-
425
- boxes = boxes[keep_idx]
426
- cls_ids = cls_ids[keep_idx]
427
-
428
- pad_w, pad_h = pad
429
- orig_w, orig_h = orig_size
430
-
431
- boxes[:, [0, 2]] -= pad_w
432
- boxes[:, [1, 3]] -= pad_h
433
- boxes /= ratio
434
- boxes = self._clip_boxes(boxes, (orig_w, orig_h))
435
 
436
  results: list[BoundingBox] = []
437
- for box, conf, cls_id in zip(boxes, scores, cls_ids):
438
- x1, y1, x2, y2 = box.tolist()
439
-
440
  if x2 <= x1 or y2 <= y1:
441
  continue
442
-
443
  results.append(
444
  BoundingBox(
445
  x1=int(math.floor(x1)),
446
  y1=int(math.floor(y1)),
447
  x2=int(math.ceil(x2)),
448
  y2=int(math.ceil(y2)),
449
- cls_id=int(cls_id),
450
  conf=float(conf),
451
  )
452
  )
453
-
454
  return results
455
 
456
- def _postprocess(
457
- self,
458
- output: np.ndarray,
459
- ratio: float,
460
- pad: tuple[float, float],
461
- orig_size: tuple[int, int],
462
- ) -> list[BoundingBox]:
463
- """
464
- Prefer final detections first.
465
- Fallback to raw decode only if needed.
466
- """
467
- # final detections: [N,6]
468
- if output.ndim == 2 and output.shape[1] >= 6:
469
- return self._decode_final_dets(output, ratio, pad, orig_size)
470
-
471
- # final detections: [1,N,6]
472
- if output.ndim == 3 and output.shape[0] == 1 and output.shape[2] == 6:
473
- return self._decode_final_dets(output, ratio, pad, orig_size)
474
-
475
- # fallback raw decode
476
- return self._decode_raw_yolo(output, ratio, pad, orig_size)
477
-
478
- def _predict_single(self, image: np.ndarray) -> list[BoundingBox]:
479
- if image is None:
480
- raise ValueError("Input image is None")
481
- if not isinstance(image, np.ndarray):
482
- raise TypeError(f"Input is not numpy array: {type(image)}")
483
- if image.ndim != 3:
484
- raise ValueError(f"Expected HWC image, got shape={image.shape}")
485
- if image.shape[0] <= 0 or image.shape[1] <= 0:
486
- raise ValueError(f"Invalid image shape={image.shape}")
487
- if image.shape[2] != 3:
488
- raise ValueError(f"Expected 3 channels, got shape={image.shape}")
489
-
490
- if image.dtype != np.uint8:
491
- image = image.astype(np.uint8)
492
-
493
- input_tensor, ratio, pad, orig_size = self._preprocess(image)
494
-
495
- expected_shape = (1, 3, self.input_height, self.input_width)
496
- if input_tensor.shape != expected_shape:
497
- raise ValueError(
498
- f"Bad input tensor shape={input_tensor.shape}, expected={expected_shape}"
499
- )
500
-
501
- outputs = self.session.run(self.output_names, {self.input_name: input_tensor})
502
- det_output = outputs[0]
503
- return self._postprocess(det_output, ratio, pad, orig_size)
504
-
505
- def _predict_tta(self, image: np.ndarray) -> list[BoundingBox]:
506
- """Horizontal-flip TTA: merge original + flipped via hard NMS."""
507
- boxes_orig = self._predict_single(image)
508
-
509
- flipped = cv2.flip(image, 1)
510
- boxes_flip = self._predict_single(flipped)
511
-
512
- w = image.shape[1]
513
- boxes_flip = [
514
- BoundingBox(
515
- x1=w - b.x2, y1=b.y1, x2=w - b.x1, y2=b.y2,
516
- cls_id=b.cls_id, conf=b.conf,
517
- )
518
- for b in boxes_flip
519
- ]
520
-
521
- all_boxes = boxes_orig + boxes_flip
522
- if len(all_boxes) == 0:
523
- return []
524
-
525
- coords = np.array(
526
- [[b.x1, b.y1, b.x2, b.y2] for b in all_boxes], dtype=np.float32
527
- )
528
- scores = np.array([b.conf for b in all_boxes], dtype=np.float32)
529
-
530
- hard_keep = self._hard_nms(coords, scores, self.iou_thres)
531
- if len(hard_keep) == 0:
532
- return []
533
-
534
- # _hard_nms already orders kept indices by descending score.
535
- hard_keep = hard_keep[: self.max_det]
536
-
537
- return [
538
- BoundingBox(
539
- x1=all_boxes[i].x1,
540
- y1=all_boxes[i].y1,
541
- x2=all_boxes[i].x2,
542
- y2=all_boxes[i].y2,
543
- cls_id=all_boxes[i].cls_id,
544
- conf=float(scores[i]),
545
- )
546
- for i in hard_keep
547
- ]
548
-
549
  def predict_batch(
550
  self,
551
  batch_images: list[ndarray],
@@ -553,24 +418,12 @@ class Miner:
553
  n_keypoints: int,
554
  ) -> list[TVFrameResult]:
555
  results: list[TVFrameResult] = []
556
-
557
  for frame_number_in_batch, image in enumerate(batch_images):
558
  try:
559
- if self.use_tta:
560
- boxes = self._predict_tta(image)
561
- else:
562
- boxes = self._predict_single(image)
563
  except Exception as e:
564
- print(f"⚠️ Inference failed for frame {offset + frame_number_in_batch}: {e}")
565
  boxes = []
566
- # for box in boxes:
567
- # if box.cls_id == 2:
568
- # box.cls_id = 3
569
- # elif box.cls_id == 3:
570
- # box.cls_id = 2
571
-
572
-
573
-
574
  results.append(
575
  TVFrameResult(
576
  frame_id=offset + frame_number_in_batch,
@@ -578,55 +431,4 @@ class Miner:
578
  keypoints=[(0, 0) for _ in range(max(0, int(n_keypoints)))],
579
  )
580
  )
581
-
582
  return results
583
-
584
-
585
- if __name__ == "__main__":
586
- # Simple manual test: load weights.onnx, run on 1.png, and draw bboxes
587
- repo_dir = Path(__file__).parent
588
- miner = Miner(repo_dir)
589
-
590
- image_path = repo_dir / "car1.png"
591
- if not image_path.exists():
592
- raise FileNotFoundError(f"Test image not found: {image_path}")
593
-
594
- image = cv2.imread(str(image_path), cv2.IMREAD_COLOR)
595
- if image is None:
596
- raise RuntimeError(f"Failed to read image: {image_path}")
597
-
598
- results = miner.predict_batch([image], offset=0, n_keypoints=0)
599
- # Draw bounding boxes on a copy of the image
600
- vis = image.copy()
601
- colors = [(0, 255, 0), (0, 0, 255), (255, 0, 0)]
602
- for frame in results:
603
- print(f"Frame {frame.frame_id}:")
604
- for i, box in enumerate(frame.boxes):
605
- color = colors[i % len(colors)]
606
- cv2.rectangle(
607
- vis,
608
- (box.x1, box.y1),
609
- (box.x2, box.y2),
610
- color,
611
- 2,
612
- )
613
- label = f"{box.cls_id }_{miner.class_names[box.cls_id] if box.cls_id < len(miner.class_names) else box.cls_id}:{box.conf:.2f}"
614
- cv2.putText(
615
- vis,
616
- label,
617
- (box.x1, max(0, box.y1 - 5)),
618
- cv2.FONT_HERSHEY_SIMPLEX,
619
- box.conf,
620
- color,
621
- 1,
622
- cv2.LINE_AA,
623
- )
624
- print(
625
- f" cls={box.cls_id} conf={box.conf:.3f} "
626
- f"box=({box.x1},{box.y1},{box.x2},{box.y2})"
627
- )
628
- print(len(frame.boxes))
629
-
630
- out_path = repo_dir / f"1_out_iou{miner.iou_thres:.2f}.png"
631
- cv2.imwrite(str(out_path), vis)
632
- print(f"Saved visualization to: {out_path}")
 
1
+ """Plate-detection miner — v2.0 "hermestech + tile s<3".
2
+
3
+ Base weights: hermestech00/numberplate0 (YOLO26s retrained, fp16, ~19 MB).
4
+
5
+ Inference pipeline:
6
+ 1) Full-image primary pass with alfred001 tuning
7
+ (conf=0.22, iou=0.41, sigma=0.685, soft-NMS + hflip TTA).
8
+ 2) If the primary returned fewer than 3 boxes, run a 2x2
9
+ overlapping tile pass (tile_conf=0.40) with novelty-merge.
10
+ """
11
  from pathlib import Path
12
  import math
13
 
 
32
  boxes: list[BoundingBox]
33
  keypoints: list[tuple[int, int]]
34
 
35
+
36
  SIZE = 1280
37
 
38
 
 
48
  if ln.strip() and not ln.strip().startswith("#")
49
  ]
50
  else:
51
+ self.class_names = ["numberplate"]
52
  print("ORT version:", ort.__version__)
53
 
54
  try:
55
  ort.preload_dlls()
56
+ print("onnxruntime.preload_dlls() success")
57
  except Exception as e:
58
+ print(f"preload_dlls failed: {e}")
59
 
60
  print("ORT available providers BEFORE session:", ort.get_available_providers())
61
 
62
+ try:
63
+ import torch
64
+ if torch.cuda.is_available():
65
+ print(f"GPU: {torch.cuda.get_device_name(0)}")
66
+ print(f"GPU memory: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f} GB")
67
+ else:
68
+ print("GPU: CUDA not available via torch")
69
+ except Exception as e:
70
+ print(f"GPU detection failed: {e}")
71
+
72
  sess_options = ort.SessionOptions()
73
  sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
74
 
 
78
  sess_options=sess_options,
79
  providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
80
  )
81
+ print("Created ORT session with preferred CUDA provider list")
82
  except Exception as e:
83
+ print(f"CUDA session creation failed, falling back to CPU: {e}")
84
  self.session = ort.InferenceSession(
85
  str(model_path),
86
  sess_options=sess_options,
 
91
 
92
  for inp in self.session.get_inputs():
93
  print("INPUT:", inp.name, inp.shape, inp.type)
 
94
  for out in self.session.get_outputs():
95
  print("OUTPUT:", out.name, out.shape, out.type)
96
 
97
  self.input_name = self.session.get_inputs()[0].name
98
+ self.output_names = [o.name for o in self.session.get_outputs()]
99
  self.input_shape = self.session.get_inputs()[0].shape
100
 
101
  self.input_height = self._safe_dim(self.input_shape[2], default=SIZE)
102
  self.input_width = self._safe_dim(self.input_shape[3], default=SIZE)
103
 
104
+ # Primary pass: alfred001 tuning (optimized for hermestech weights)
105
+ self.conf_thres = 0.22
106
+ self.iou_thres = 0.41
107
+ self.sigma = 0.685
108
+ self.max_det = 300
109
+
110
+ # Conditional tile-pass (trimmed for latency: no hflip, tighter sparse)
111
+ self.sparse_threshold = 3 # fire tiles only if primary returns < this
112
+ self.tile_conf = 0.40
113
+ self.tile_overlap = 0.20
114
+ self.novelty_iou = 0.10
115
+ self.final_max_det = 22
116
+ self.tile_use_hflip = False # skip hflip tile pass to save ~4 forwards
117
+
118
  self.use_tta = True
119
 
120
+ print(f"ONNX model loaded from: {model_path}")
121
+ print(f"ONNX providers: {self.session.get_providers()}")
122
+ print(f"ONNX input: name={self.input_name}, shape={self.input_shape}")
123
 
124
  def __repr__(self) -> str:
125
  return (
 
131
  def _safe_dim(value, default: int) -> int:
132
  return value if isinstance(value, int) and value > 0 else default
133
 
134
+ # ---------- image preprocessing ----------
135
  def _letterbox(
136
  self,
137
  image: ndarray,
138
  new_shape: tuple[int, int],
139
  color=(114, 114, 114),
140
  ) -> tuple[ndarray, float, tuple[float, float]]:
 
 
 
 
 
 
 
141
  h, w = image.shape[:2]
142
  new_w, new_h = new_shape
 
143
  ratio = min(new_w / w, new_h / h)
144
  resized_w = int(round(w * ratio))
145
  resized_h = int(round(h * ratio))
 
146
  if (resized_w, resized_h) != (w, h):
147
  interp = cv2.INTER_CUBIC if ratio > 1.0 else cv2.INTER_LINEAR
148
  image = cv2.resize(image, (resized_w, resized_h), interpolation=interp)
149
+ dw = (new_w - resized_w) / 2.0
150
+ dh = (new_h - resized_h) / 2.0
 
 
 
 
151
  left = int(round(dw - 0.1))
152
  right = int(round(dw + 0.1))
153
  top = int(round(dh - 0.1))
154
  bottom = int(round(dh + 0.1))
 
155
  padded = cv2.copyMakeBorder(
156
+ image, top, bottom, left, right,
157
+ borderType=cv2.BORDER_CONSTANT, value=color,
 
 
 
 
 
158
  )
159
  return padded, ratio, (dw, dh)
160
 
161
+ def _preprocess(self, image: ndarray):
162
+ img, ratio, pad = self._letterbox(image, (self.input_width, self.input_height))
163
+ img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
164
  img = np.transpose(img, (2, 0, 1))[None, ...]
165
+ return np.ascontiguousarray(img, dtype=np.float32), ratio, pad
 
 
166
 
167
  @staticmethod
168
  def _clip_boxes(boxes: np.ndarray, image_size: tuple[int, int]) -> np.ndarray:
 
173
  boxes[:, 3] = np.clip(boxes[:, 3], 0, h - 1)
174
  return boxes
175
 
176
+ # ---------- NMS primitives ----------
177
  @staticmethod
178
+ def _hard_nms(boxes: np.ndarray, scores: np.ndarray, iou_thresh: float) -> np.ndarray:
179
+ N = len(boxes)
180
+ if N == 0:
181
+ return np.array([], dtype=np.intp)
182
+ boxes = np.asarray(boxes, dtype=np.float32)
183
+ scores = np.asarray(scores, dtype=np.float32)
184
+ order = np.argsort(-scores)
185
+ keep: list[int] = []
186
+ while len(order):
187
+ i = int(order[0])
188
+ keep.append(i)
189
+ if len(order) == 1:
190
+ break
191
+ rest = order[1:]
192
+ xx1 = np.maximum(boxes[i, 0], boxes[rest, 0])
193
+ yy1 = np.maximum(boxes[i, 1], boxes[rest, 1])
194
+ xx2 = np.minimum(boxes[i, 2], boxes[rest, 2])
195
+ yy2 = np.minimum(boxes[i, 3], boxes[rest, 3])
196
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
197
+ area_i = (boxes[i, 2] - boxes[i, 0]) * (boxes[i, 3] - boxes[i, 1])
198
+ area_r = (boxes[rest, 2] - boxes[rest, 0]) * (boxes[rest, 3] - boxes[rest, 1])
199
+ iou = inter / (area_i + area_r - inter + 1e-7)
200
+ order = rest[iou <= iou_thresh]
201
+ return np.array(keep, dtype=np.intp)
202
 
203
  def _soft_nms(
204
  self,
205
  boxes: np.ndarray,
206
  scores: np.ndarray,
207
+ sigma: float,
208
  score_thresh: float = 0.01,
209
  ) -> tuple[np.ndarray, np.ndarray]:
 
 
 
 
210
  N = len(boxes)
211
  if N == 0:
212
  return np.array([], dtype=np.intp), np.array([], dtype=np.float32)
 
213
  boxes = boxes.astype(np.float32, copy=True)
214
  scores = scores.astype(np.float32, copy=True)
215
  order = np.arange(N)
 
216
  for i in range(N):
217
  max_pos = i + int(np.argmax(scores[i:]))
218
  boxes[[i, max_pos]] = boxes[[max_pos, i]]
219
  scores[[i, max_pos]] = scores[[max_pos, i]]
220
  order[[i, max_pos]] = order[[max_pos, i]]
 
221
  if i + 1 >= N:
222
  break
 
223
  xx1 = np.maximum(boxes[i, 0], boxes[i + 1:, 0])
224
  yy1 = np.maximum(boxes[i, 1], boxes[i + 1:, 1])
225
  xx2 = np.minimum(boxes[i, 2], boxes[i + 1:, 2])
226
  yy2 = np.minimum(boxes[i, 3], boxes[i + 1:, 3])
227
  inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
228
+ area_i = float(
 
229
  (boxes[i, 2] - boxes[i, 0]) * (boxes[i, 3] - boxes[i, 1])
230
+ )
231
  areas_j = (
232
  np.maximum(0.0, boxes[i + 1:, 2] - boxes[i + 1:, 0])
233
  * np.maximum(0.0, boxes[i + 1:, 3] - boxes[i + 1:, 1])
234
  )
235
  iou = inter / (area_i + areas_j - inter + 1e-7)
236
  scores[i + 1:] *= np.exp(-(iou ** 2) / sigma)
 
237
  mask = scores > score_thresh
238
  return order[mask], scores[mask]
239
 
240
  @staticmethod
241
+ def _box_iou_one_to_many(box: np.ndarray, boxes: np.ndarray) -> np.ndarray:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
242
  if len(boxes) == 0:
243
+ return np.zeros(0, dtype=np.float32)
244
+ xx1 = np.maximum(box[0], boxes[:, 0])
245
+ yy1 = np.maximum(box[1], boxes[:, 1])
246
+ xx2 = np.minimum(box[2], boxes[:, 2])
247
+ yy2 = np.minimum(box[3], boxes[:, 3])
248
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
249
+ area_a = max(0.0, (box[2] - box[0]) * (box[3] - box[1]))
250
+ area_b = np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) * np.maximum(0.0, boxes[:, 3] - boxes[:, 1])
251
+ return inter / (area_a + area_b - inter + 1e-7)
252
+
253
+ # ---------- raw-dets helper ----------
254
+ def _raw_dets(self, image: ndarray, conf: float) -> np.ndarray:
255
+ """Run a single forward pass and return [N, 5] dets in ORIGINAL image coords."""
256
+ x, ratio, (dw, dh) = self._preprocess(image)
257
+ out = self.session.run(self.output_names, {self.input_name: x})[0]
258
+ if out.ndim == 3:
259
+ out = out[0]
260
+ if out.shape[1] < 5:
261
+ return np.zeros((0, 5), dtype=np.float32)
262
+ boxes = out[:, :4].astype(np.float32)
263
+ scores = out[:, 4].astype(np.float32)
264
+ keep = scores >= conf
265
+ boxes, scores = boxes[keep], scores[keep]
266
+ if len(boxes) == 0:
267
+ return np.zeros((0, 5), dtype=np.float32)
268
+ boxes[:, [0, 2]] -= dw
269
+ boxes[:, [1, 3]] -= dh
270
  boxes /= ratio
271
+ oh, ow = image.shape[:2]
272
+ boxes = self._clip_boxes(boxes, (ow, oh))
273
+ return np.concatenate([boxes, scores[:, None]], axis=1)
 
 
 
 
 
 
 
274
 
275
+ # ---------- primary pass: soft-NMS + hflip TTA ----------
276
+ def _primary(self, image: ndarray) -> np.ndarray:
277
+ d1 = self._raw_dets(image, self.conf_thres)
278
+ flipped = cv2.flip(image, 1)
279
+ d2 = self._raw_dets(flipped, self.conf_thres)
280
+ if len(d2):
281
+ w = image.shape[1]
282
+ x1 = w - d2[:, 2]
283
+ x2 = w - d2[:, 0]
284
+ d2 = np.stack([x1, d2[:, 1], x2, d2[:, 3], d2[:, 4]], axis=1)
285
+ all_d = np.concatenate([d1, d2], axis=0) if len(d2) else d1
286
+ if len(all_d) == 0:
287
+ return np.zeros((0, 5), dtype=np.float32)
288
+ # soft-NMS, then hard-NMS
289
+ keep_idx, scores = self._soft_nms(all_d[:, :4].copy(), all_d[:, 4].copy(), sigma=self.sigma)
290
+ if len(keep_idx) == 0:
291
+ return np.zeros((0, 5), dtype=np.float32)
292
+ merged = np.concatenate([all_d[keep_idx, :4], scores[:, None]], axis=1)
293
+ keep = self._hard_nms(merged[:, :4], merged[:, 4], self.iou_thres)
294
+ merged = merged[keep]
295
+ if len(merged) > self.max_det:
296
+ merged = merged[np.argsort(-merged[:, 4])[: self.max_det]]
297
+ return merged
298
+
299
+ # ---------- conditional tile pass ----------
300
+ def _tile_augment(self, image: ndarray, primary: np.ndarray) -> np.ndarray:
301
+ """Run 2x2 overlapping tiles + hflip, novelty-merge into primary."""
302
+ oh, ow = image.shape[:2]
303
+ tw, th = ow // 2, oh // 2
304
+ ox, oy = int(tw * self.tile_overlap), int(th * self.tile_overlap)
305
+ tiles = [
306
+ (0, 0, min(ow, tw + ox), min(oh, th + oy)),
307
+ (max(0, tw - ox), 0, ow, min(oh, th + oy)),
308
+ (0, max(0, th - oy), min(ow, tw + ox), oh),
309
+ (max(0, tw - ox), max(0, th - oy), ow, oh),
310
+ ]
311
+ collected: list[np.ndarray] = []
312
+ for x1, y1, x2, y2 in tiles:
313
+ crop = image[y1:y2, x1:x2]
314
+ if crop.size == 0:
315
  continue
316
+ d = self._raw_dets(crop, self.tile_conf)
317
+ if len(d):
318
+ d[:, 0] += x1
319
+ d[:, 1] += y1
320
+ d[:, 2] += x1
321
+ d[:, 3] += y1
322
+ collected.append(d)
323
+
324
+ # hflip tile pass (skipped when tile_use_hflip=False — saves 4 ONNX forwards)
325
+ if self.tile_use_hflip:
326
+ flipped = cv2.flip(image, 1)
327
+ for x1, y1, x2, y2 in tiles:
328
+ fx1 = ow - x2
329
+ fx2 = ow - x1
330
+ if fx2 <= fx1:
331
+ continue
332
+ crop = flipped[y1:y2, fx1:fx2]
333
+ if crop.size == 0:
334
+ continue
335
+ d = self._raw_dets(crop, self.tile_conf)
336
+ if len(d):
337
+ d_un = d.copy()
338
+ d_un[:, 0] = (ow - (d[:, 2] + fx1))
339
+ d_un[:, 2] = (ow - (d[:, 0] + fx1))
340
+ d_un[:, 1] = d[:, 1] + y1
341
+ d_un[:, 3] = d[:, 3] + y1
342
+ collected.append(d_un)
343
+
344
+ if not collected:
345
+ return primary
346
+
347
+ tile_dets = np.concatenate(collected, axis=0)
348
+ keep = self._hard_nms(tile_dets[:, :4], tile_dets[:, 4], 0.5)
349
+ tile_dets = tile_dets[keep]
350
+
351
+ # Novelty: drop tile boxes that overlap any primary box at IoU >= novelty_iou
352
+ if len(primary) > 0 and len(tile_dets) > 0:
353
+ mask = np.ones(len(tile_dets), dtype=bool)
354
+ for i in range(len(tile_dets)):
355
+ ious = self._box_iou_one_to_many(tile_dets[i, :4], primary[:, :4])
356
+ if len(ious) and np.max(ious) >= self.novelty_iou:
357
+ mask[i] = False
358
+ tile_dets = tile_dets[mask]
359
+
360
+ if len(tile_dets) == 0:
361
+ return primary
362
+
363
+ # Sanity filter: min/max size, aspect ratio
364
+ w = tile_dets[:, 2] - tile_dets[:, 0]
365
+ h = tile_dets[:, 3] - tile_dets[:, 1]
366
+ area = w * h
367
+ ar = np.maximum(w / np.maximum(h, 1e-6), h / np.maximum(w, 1e-6))
368
+ img_area = float(ow * oh)
369
+ ok = (w >= 6) & (h >= 6) & (area >= 36) & (area <= 0.5 * img_area) & (ar <= 10.0)
370
+ tile_dets = tile_dets[ok]
371
+ if len(tile_dets) == 0:
372
+ return primary
373
+
374
+ merged = np.concatenate([primary, tile_dets], axis=0)
375
+ keep = self._hard_nms(merged[:, :4], merged[:, 4], self.iou_thres)
376
+ merged = merged[keep]
377
+ if len(merged) > self.final_max_det:
378
+ merged = merged[np.argsort(-merged[:, 4])[: self.final_max_det]]
379
+ return merged
380
+
381
+ # ---------- single-image predict ----------
382
+ def _predict_single(self, image: ndarray) -> list[BoundingBox]:
383
+ if image is None or not isinstance(image, np.ndarray) or image.ndim != 3:
384
  return []
385
+ if image.shape[0] <= 0 or image.shape[1] <= 0 or image.shape[2] != 3:
386
+ return []
387
+ if image.dtype != np.uint8:
388
+ image = image.astype(np.uint8)
389
 
390
+ primary = self._primary(image)
391
+ if len(primary) < self.sparse_threshold:
392
+ dets = self._tile_augment(image, primary)
393
+ else:
394
+ dets = primary
 
 
 
 
 
 
 
 
 
 
395
 
396
  results: list[BoundingBox] = []
397
+ for row in dets:
398
+ x1, y1, x2, y2, conf = row.tolist()
 
399
  if x2 <= x1 or y2 <= y1:
400
  continue
 
401
  results.append(
402
  BoundingBox(
403
  x1=int(math.floor(x1)),
404
  y1=int(math.floor(y1)),
405
  x2=int(math.ceil(x2)),
406
  y2=int(math.ceil(y2)),
407
+ cls_id=0,
408
  conf=float(conf),
409
  )
410
  )
 
411
  return results
412
 
413
+ # ---------- chute entrypoint ----------
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
414
  def predict_batch(
415
  self,
416
  batch_images: list[ndarray],
 
418
  n_keypoints: int,
419
  ) -> list[TVFrameResult]:
420
  results: list[TVFrameResult] = []
 
421
  for frame_number_in_batch, image in enumerate(batch_images):
422
  try:
423
+ boxes = self._predict_single(image)
 
 
 
424
  except Exception as e:
425
+ print(f"Inference failed for frame {offset + frame_number_in_batch}: {e}")
426
  boxes = []
 
 
 
 
 
 
 
 
427
  results.append(
428
  TVFrameResult(
429
  frame_id=offset + frame_number_in_batch,
 
431
  keypoints=[(0, 0) for _ in range(max(0, int(n_keypoints)))],
432
  )
433
  )
 
434
  return results