thomaskk2024 commited on
Commit
cd3b5ca
·
verified ·
1 Parent(s): 2a6d118

Upload folder using huggingface_hub

Browse files
Files changed (3) hide show
  1. chute_config.yml +21 -0
  2. miner.py +866 -0
  3. weights.onnx +3 -0
chute_config.yml ADDED
@@ -0,0 +1,21 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ Image:
2
+ from_base: parachutes/python:3.12
3
+ run_command:
4
+ - pip install --upgrade setuptools wheel
5
+ - pip install 'numpy>=1.23' 'onnxruntime-gpu[cuda,cudnn]>=1.16' 'opencv-python>=4.7' 'pillow>=9.5' 'huggingface_hub>=0.19.4' 'pydantic>=2.0' 'pyyaml>=6.0' 'aiohttp>=3.9'
6
+ - pip install torch torchvision
7
+
8
+ NodeSelector:
9
+ gpu_count: 1
10
+ min_vram_gb_per_gpu: 16
11
+ max_hourly_price_per_gpu: 2
12
+ include:
13
+ - pro_6000
14
+
15
+ Chute:
16
+ timeout_seconds: 900
17
+ concurrency: 4
18
+ max_instances: 5
19
+ scaling_threshold: 0.5
20
+ shutdown_after_seconds: 288000
21
+ tee: true
miner.py ADDED
@@ -0,0 +1,866 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from pathlib import Path
2
+ import math
3
+
4
+ import cv2
5
+ import numpy as np
6
+ import onnxruntime as ort
7
+ from numpy import ndarray
8
+ from pydantic import BaseModel
9
+
10
+
11
+ class BoundingBox(BaseModel):
12
+ x1: int
13
+ y1: int
14
+ x2: int
15
+ y2: int
16
+ cls_id: int
17
+ conf: float
18
+
19
+
20
+ class TVFrameResult(BaseModel):
21
+ frame_id: int
22
+ boxes: list[BoundingBox]
23
+ keypoints: list[tuple[int, int]]
24
+
25
+
26
+ class Miner:
27
+ """ONNX Runtime miner for road-sign detection (single class).
28
+
29
+ Strategy (ported from offense / fire001 miner):
30
+ - per-class confidence threshold with per-class rescue bonus
31
+ - per-class hard NMS, then cross-class dedup (no-op for single class)
32
+ - horizontal-flip TTA with full-set cluster score boost
33
+ Plus: class remap, sanity-box filter tuned for small distant signs,
34
+ TTA toggle.
35
+ """
36
+
37
+ class_names = ["road_sign"]
38
+ # Order the model emits classes in -- remapped to `class_names` index.
39
+ _model_class_order = ["road_sign"]
40
+
41
+ iou_thres = 0.5
42
+ cross_iou_thresh = 0.8
43
+ max_det = 150
44
+
45
+ # Per-class confidence threshold. Road signs in this dataset are
46
+ # frequently degraded / rear-facing / partly-obscured / distant, so we
47
+ # run noticeably below the fire/smoke baseline. The validator's
48
+ # false_positive pillar = max(0, 1 - ffpi/10): we can tolerate ~2 FP per
49
+ # image and still keep that pillar above 0.8.
50
+ _conf_thres_array = np.array(
51
+ [0.35], dtype=np.float32
52
+ )
53
+ # Per-class rescue bonus. If a class has ZERO boxes passing the threshold
54
+ # in a frame, its top-1 candidate is admitted when its score is at least
55
+ # (threshold - bonus). Bumped from 0.05 -> 0.10 so a single faint sign in
56
+ # an otherwise empty frame still produces a detection (map50 recall win,
57
+ # at most one extra FP per such frame).
58
+ _bonus_array = np.array(
59
+ [0.2], dtype=np.float32
60
+ )
61
+
62
+ # Box sanity filter: drop tiny / degenerate / image-spanning / extreme
63
+ # aspect ratio boxes.
64
+ # min_box_area = 14x14 -> 14x14 is the smallest credible sign. The old
65
+ # value of 64 (8x8) silently discarded narrow
66
+ # distant signs like a 10x6 px overhead chevron.
67
+ # min_side = 3 -> matches min_box_area; anything thinner is
68
+ # almost certainly a pole or shadow false alarm.
69
+ # max_aspect_ratio = 12.0
70
+ # -> overhead destination panels and lane-assignment
71
+ # signs are very wide (long, thin rectangles);
72
+ # 8.0 was clipping legitimate detections.
73
+ min_box_area = 14 * 14
74
+ min_side = 3
75
+ max_aspect_ratio = 12.0
76
+
77
+ # Tile-based TTA: when the source image is significantly larger than the
78
+ # model input, letterboxing throws away ~1.5x of effective resolution,
79
+ # which kills small-sign recall. Splitting into overlapping horizontal
80
+ # tiles preserves native resolution on each half. Triggered only when
81
+ # source width >= tile_trigger_ratio * model_input_width to avoid wasted
82
+ # compute on already-small images.
83
+ tile_trigger_ratio = 1.4
84
+ tile_overlap_ratio = 0.20
85
+
86
+ def __init__(self, path_hf_repo: Path) -> None:
87
+ model_path = path_hf_repo / "weights.onnx"
88
+ self.cls_remap = np.array(
89
+ [self.class_names.index(n) for n in self._model_class_order],
90
+ dtype=np.int32,
91
+ )
92
+ print("ORT version:", ort.__version__)
93
+
94
+ try:
95
+ ort.preload_dlls()
96
+ print("✅ onnxruntime.preload_dlls() success")
97
+ except Exception as e:
98
+ print(f"⚠️ preload_dlls failed: {e}")
99
+
100
+ print("ORT available providers BEFORE session:", ort.get_available_providers())
101
+
102
+ sess_options = ort.SessionOptions()
103
+ sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
104
+ # The public-track latency gate runs CPU-only on 2 vCPUs. Without this,
105
+ # ORT spawns a thread per host core and oversubscribes the 2 cores
106
+ # (~5.5x slower). Pin to 2 intra-op threads / 1 inter-op to match.
107
+ sess_options.intra_op_num_threads = 2
108
+ sess_options.inter_op_num_threads = 1
109
+
110
+ try:
111
+ self.session = ort.InferenceSession(
112
+ str(model_path),
113
+ sess_options=sess_options,
114
+ providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
115
+ )
116
+ print("✅ Created ORT session with preferred CUDA provider list")
117
+ except Exception as e:
118
+ print(f"⚠️ CUDA session creation failed, falling back to CPU: {e}")
119
+ self.session = ort.InferenceSession(
120
+ str(model_path),
121
+ sess_options=sess_options,
122
+ providers=["CPUExecutionProvider"],
123
+ )
124
+
125
+ print("ORT session providers:", self.session.get_providers())
126
+
127
+ for inp in self.session.get_inputs():
128
+ print("INPUT:", inp.name, inp.shape, inp.type)
129
+ for out in self.session.get_outputs():
130
+ print("OUTPUT:", out.name, out.shape, out.type)
131
+
132
+ self.input_name = self.session.get_inputs()[0].name
133
+ self.output_names = [output.name for output in self.session.get_outputs()]
134
+ self.input_shape = self.session.get_inputs()[0].shape
135
+
136
+ self.input_height = self._safe_dim(self.input_shape[2], default=1024)
137
+ self.input_width = self._safe_dim(self.input_shape[3], default=1024)
138
+
139
+ self.use_tta = False
140
+ self.use_tile_tta = False
141
+ # Soft-NMS (ported from carwash001): Gaussian score decay of overlapping
142
+ # boxes instead of hard removal. OFF by default to preserve the current
143
+ # deployed behaviour; flip on (and tune sigma) via tune_miner.py to see if
144
+ # it scores better — useful where signs cluster (gantries, sign assemblies).
145
+ self.use_soft_nms = False
146
+ self.soft_nms_sigma = 0.5
147
+ self.soft_nms_score_thresh = 0.01
148
+
149
+ print(f"✅ ONNX model loaded from: {model_path}")
150
+ print(f"✅ ONNX providers: {self.session.get_providers()}")
151
+ print(f"✅ ONNX input: name={self.input_name}, shape={self.input_shape}")
152
+ print("per-class conf: " + ", ".join(
153
+ f"{n}={t:.3f}" for n, t in zip(
154
+ self.class_names, self._conf_thres_array.tolist()
155
+ )
156
+ ))
157
+
158
+ def __repr__(self) -> str:
159
+ return (
160
+ f"ONNXRuntime(session={type(self.session).__name__}, "
161
+ f"providers={self.session.get_providers()})"
162
+ )
163
+
164
+ @staticmethod
165
+ def _safe_dim(value, default: int) -> int:
166
+ return value if isinstance(value, int) and value > 0 else default
167
+
168
+ def _letterbox(
169
+ self,
170
+ image: ndarray,
171
+ new_shape: tuple[int, int],
172
+ color=(114, 114, 114),
173
+ ) -> tuple[ndarray, float, tuple[float, float]]:
174
+ h, w = image.shape[:2]
175
+ new_w, new_h = new_shape
176
+
177
+ ratio = min(new_w / w, new_h / h)
178
+ resized_w = int(round(w * ratio))
179
+ resized_h = int(round(h * ratio))
180
+
181
+ if (resized_w, resized_h) != (w, h):
182
+ interp = cv2.INTER_CUBIC if ratio > 1.0 else cv2.INTER_LINEAR
183
+ image = cv2.resize(image, (resized_w, resized_h), interpolation=interp)
184
+
185
+ dw = (new_w - resized_w) / 2.0
186
+ dh = (new_h - resized_h) / 2.0
187
+
188
+ left = int(round(dw - 0.1))
189
+ right = int(round(dw + 0.1))
190
+ top = int(round(dh - 0.1))
191
+ bottom = int(round(dh + 0.1))
192
+
193
+ padded = cv2.copyMakeBorder(
194
+ image, top, bottom, left, right,
195
+ borderType=cv2.BORDER_CONSTANT, value=color,
196
+ )
197
+ return padded, ratio, (dw, dh)
198
+
199
+ def _preprocess(
200
+ self, image: ndarray
201
+ ) -> tuple[np.ndarray, float, tuple[float, float], tuple[int, int]]:
202
+ orig_h, orig_w = image.shape[:2]
203
+ img, ratio, pad = self._letterbox(
204
+ image, (self.input_width, self.input_height)
205
+ )
206
+ img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
207
+ img = img.astype(np.float32) / 255.0
208
+ img = np.transpose(img, (2, 0, 1))[None, ...]
209
+ img = np.ascontiguousarray(img, dtype=np.float32)
210
+ return img, ratio, pad, (orig_w, orig_h)
211
+
212
+ @staticmethod
213
+ def _clip_boxes(boxes: np.ndarray, image_size: tuple[int, int]) -> np.ndarray:
214
+ w, h = image_size
215
+ boxes[:, 0] = np.clip(boxes[:, 0], 0, w - 1)
216
+ boxes[:, 1] = np.clip(boxes[:, 1], 0, h - 1)
217
+ boxes[:, 2] = np.clip(boxes[:, 2], 0, w - 1)
218
+ boxes[:, 3] = np.clip(boxes[:, 3], 0, h - 1)
219
+ return boxes
220
+
221
+ @staticmethod
222
+ def _xywh_to_xyxy(boxes: np.ndarray) -> np.ndarray:
223
+ out = np.empty_like(boxes)
224
+ out[:, 0] = boxes[:, 0] - boxes[:, 2] / 2.0
225
+ out[:, 1] = boxes[:, 1] - boxes[:, 3] / 2.0
226
+ out[:, 2] = boxes[:, 0] + boxes[:, 2] / 2.0
227
+ out[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.0
228
+ return out
229
+
230
+ @staticmethod
231
+ def _hard_nms(
232
+ boxes: np.ndarray, scores: np.ndarray, iou_thresh: float
233
+ ) -> np.ndarray:
234
+ n = len(boxes)
235
+ if n == 0:
236
+ return np.array([], dtype=np.intp)
237
+ order = np.argsort(-scores)
238
+ keep: list[int] = []
239
+ while len(order) > 0:
240
+ i = int(order[0])
241
+ keep.append(i)
242
+ if len(order) == 1:
243
+ break
244
+ rest = order[1:]
245
+ xx1 = np.maximum(boxes[i, 0], boxes[rest, 0])
246
+ yy1 = np.maximum(boxes[i, 1], boxes[rest, 1])
247
+ xx2 = np.minimum(boxes[i, 2], boxes[rest, 2])
248
+ yy2 = np.minimum(boxes[i, 3], boxes[rest, 3])
249
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
250
+ a_i = (max(0.0, boxes[i, 2] - boxes[i, 0]) *
251
+ max(0.0, boxes[i, 3] - boxes[i, 1]))
252
+ a_r = (np.maximum(0.0, boxes[rest, 2] - boxes[rest, 0]) *
253
+ np.maximum(0.0, boxes[rest, 3] - boxes[rest, 1]))
254
+ iou = inter / (a_i + a_r - inter + 1e-7)
255
+ order = rest[iou <= iou_thresh]
256
+ return np.array(keep, dtype=np.intp)
257
+
258
+ def _per_class_hard_nms(
259
+ self,
260
+ boxes: np.ndarray,
261
+ scores: np.ndarray,
262
+ cls_ids: np.ndarray,
263
+ iou_thresh: float,
264
+ ) -> np.ndarray:
265
+ if len(boxes) == 0:
266
+ return np.array([], dtype=np.intp)
267
+ all_keep: list[int] = []
268
+ for c in np.unique(cls_ids):
269
+ mask = cls_ids == c
270
+ indices = np.where(mask)[0]
271
+ keep = self._hard_nms(boxes[mask], scores[mask], iou_thresh)
272
+ all_keep.extend(indices[keep].tolist())
273
+ all_keep.sort()
274
+ return np.array(all_keep, dtype=np.intp)
275
+
276
+ def _soft_nms(
277
+ self,
278
+ boxes: np.ndarray,
279
+ scores: np.ndarray,
280
+ sigma: float = 0.5,
281
+ score_thresh: float = 0.01,
282
+ ) -> tuple[np.ndarray, np.ndarray]:
283
+ """Soft-NMS: Gaussian decay of overlapping scores instead of hard removal.
284
+ Returns (kept_original_indices, updated_scores). (Ported from carwash001.)"""
285
+ N = len(boxes)
286
+ if N == 0:
287
+ return np.array([], dtype=np.intp), np.array([], dtype=np.float32)
288
+ boxes = boxes.astype(np.float32, copy=True)
289
+ scores = scores.astype(np.float32, copy=True)
290
+ order = np.arange(N)
291
+ for i in range(N):
292
+ max_pos = i + int(np.argmax(scores[i:]))
293
+ boxes[[i, max_pos]] = boxes[[max_pos, i]]
294
+ scores[[i, max_pos]] = scores[[max_pos, i]]
295
+ order[[i, max_pos]] = order[[max_pos, i]]
296
+ if i + 1 >= N:
297
+ break
298
+ xx1 = np.maximum(boxes[i, 0], boxes[i + 1:, 0])
299
+ yy1 = np.maximum(boxes[i, 1], boxes[i + 1:, 1])
300
+ xx2 = np.minimum(boxes[i, 2], boxes[i + 1:, 2])
301
+ yy2 = np.minimum(boxes[i, 3], boxes[i + 1:, 3])
302
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
303
+ area_i = max(0.0, float(
304
+ (boxes[i, 2] - boxes[i, 0]) * (boxes[i, 3] - boxes[i, 1])))
305
+ areas_j = (np.maximum(0.0, boxes[i + 1:, 2] - boxes[i + 1:, 0])
306
+ * np.maximum(0.0, boxes[i + 1:, 3] - boxes[i + 1:, 1]))
307
+ iou = inter / (area_i + areas_j - inter + 1e-7)
308
+ scores[i + 1:] *= np.exp(-(iou ** 2) / sigma)
309
+ mask = scores > score_thresh
310
+ return order[mask], scores[mask]
311
+
312
+ def _per_class_soft_nms(
313
+ self,
314
+ boxes: np.ndarray,
315
+ scores: np.ndarray,
316
+ cls_ids: np.ndarray,
317
+ sigma: float = 0.5,
318
+ score_thresh: float = 0.01,
319
+ ) -> tuple[np.ndarray, np.ndarray]:
320
+ """Soft-NMS applied independently per class. Returns (kept_idx, updated_scores)."""
321
+ if len(boxes) == 0:
322
+ return np.array([], dtype=np.intp), np.array([], dtype=np.float32)
323
+ all_keep: list[int] = []
324
+ all_scores: list[float] = []
325
+ for c in np.unique(cls_ids):
326
+ indices = np.where(cls_ids == c)[0]
327
+ keep, updated = self._soft_nms(boxes[indices], scores[indices],
328
+ sigma, score_thresh)
329
+ for k, s in zip(keep, updated):
330
+ all_keep.append(int(indices[k])); all_scores.append(float(s))
331
+ if not all_keep:
332
+ return np.array([], dtype=np.intp), np.array([], dtype=np.float32)
333
+ return np.array(all_keep, dtype=np.intp), np.array(all_scores, dtype=np.float32)
334
+
335
+ def _cross_class_dedup_op(
336
+ self,
337
+ boxes: np.ndarray,
338
+ scores: np.ndarray,
339
+ cls_ids: np.ndarray,
340
+ iou_thresh: float,
341
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
342
+ """Remove near-duplicate boxes across classes.
343
+
344
+ Order candidates by (score - per_class_threshold) margin, then by area;
345
+ keep the highest, suppress every other box with IoU > iou_thresh.
346
+ With a single road_sign class this is effectively a no-op, but the
347
+ method is kept so the pipeline stays compatible with the multi-class
348
+ miner template.
349
+ """
350
+ n = len(boxes)
351
+ if n <= 1:
352
+ return boxes, scores, cls_ids
353
+ boxes = np.asarray(boxes, dtype=np.float32)
354
+ scores = np.asarray(scores, dtype=np.float32)
355
+ cls_ids = np.asarray(cls_ids, dtype=np.int32)
356
+ areas = (np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) *
357
+ np.maximum(0.0, boxes[:, 3] - boxes[:, 1]))
358
+ margins = scores - self._conf_thres_array[cls_ids]
359
+ order = np.lexsort((-areas, -margins))
360
+ suppressed = np.zeros(n, dtype=bool)
361
+ keep: list[int] = []
362
+ for i in order:
363
+ if suppressed[i]:
364
+ continue
365
+ keep.append(int(i))
366
+ bi = boxes[i]
367
+ xx1 = np.maximum(bi[0], boxes[:, 0])
368
+ yy1 = np.maximum(bi[1], boxes[:, 1])
369
+ xx2 = np.minimum(bi[2], boxes[:, 2])
370
+ yy2 = np.minimum(bi[3], boxes[:, 3])
371
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
372
+ a_i = max(1e-7, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
373
+ iou = inter / (a_i + areas - inter + 1e-7)
374
+ dup = iou > iou_thresh
375
+ dup[i] = False
376
+ suppressed |= dup
377
+ keep_idx = np.array(keep, dtype=np.intp)
378
+ return boxes[keep_idx], scores[keep_idx], cls_ids[keep_idx]
379
+
380
+ @staticmethod
381
+ def _max_score_per_cluster(
382
+ post_boxes: np.ndarray,
383
+ post_cls: np.ndarray,
384
+ full_boxes: np.ndarray,
385
+ full_scores: np.ndarray,
386
+ full_cls: np.ndarray,
387
+ iou_thresh: float,
388
+ ) -> np.ndarray:
389
+ """For each kept (post-NMS) box, return the max score over the FULL
390
+ candidate set among same-class boxes with IoU >= iou_thresh.
391
+
392
+ Used after horizontal-flip TTA: a high-confidence flipped detection
393
+ can raise the score of the corresponding original detection.
394
+ """
395
+ n = len(post_boxes)
396
+ if n == 0:
397
+ return np.empty(0, dtype=np.float32)
398
+ full_areas = (np.maximum(0.0, full_boxes[:, 2] - full_boxes[:, 0]) *
399
+ np.maximum(0.0, full_boxes[:, 3] - full_boxes[:, 1]))
400
+ out = np.empty(n, dtype=np.float32)
401
+ for i in range(n):
402
+ bi = post_boxes[i]
403
+ xx1 = np.maximum(bi[0], full_boxes[:, 0])
404
+ yy1 = np.maximum(bi[1], full_boxes[:, 1])
405
+ xx2 = np.minimum(bi[2], full_boxes[:, 2])
406
+ yy2 = np.minimum(bi[3], full_boxes[:, 3])
407
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
408
+ a_i = max(0.0, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
409
+ iou = inter / (a_i + full_areas - inter + 1e-7)
410
+ cluster = (iou >= iou_thresh) & (full_cls == post_cls[i])
411
+ out[i] = float(np.max(full_scores[cluster])) if np.any(cluster) else 0.0
412
+ return out
413
+
414
+ def _conf_filter_mask(
415
+ self, scores: np.ndarray, cls_ids: np.ndarray
416
+ ) -> np.ndarray:
417
+ """Boolean keep-mask: score >= per-class threshold, with a per-class
418
+ rescue -- if a class has zero boxes passing, admit its top-1 candidate
419
+ when its score >= (per-class threshold - per-class bonus)."""
420
+ if len(scores) == 0:
421
+ return np.zeros(0, dtype=bool)
422
+ thr = self._conf_thres_array[cls_ids]
423
+ keep = scores >= thr
424
+ for c in np.unique(cls_ids):
425
+ b = float(self._bonus_array[c])
426
+ if b <= 0.0:
427
+ continue
428
+ cm = cls_ids == c
429
+ if keep[cm].any():
430
+ continue
431
+ idx = np.where(cm)[0]
432
+ top = int(idx[int(np.argmax(scores[idx]))])
433
+ if scores[top] >= self._conf_thres_array[c] - b:
434
+ keep[top] = True
435
+ return keep
436
+
437
+ def _filter_sane_boxes(
438
+ self,
439
+ boxes: np.ndarray,
440
+ scores: np.ndarray,
441
+ cls_ids: np.ndarray,
442
+ orig_size: tuple[int, int],
443
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
444
+ """Drop tiny / degenerate / image-spanning / extreme-AR boxes (FP)."""
445
+ if len(boxes) == 0:
446
+ return boxes, scores, cls_ids
447
+ orig_w, orig_h = orig_size
448
+ image_area = float(orig_w * orig_h)
449
+ keep = []
450
+ for i, box in enumerate(boxes):
451
+ x1, y1, x2, y2 = box.tolist()
452
+ bw = x2 - x1
453
+ bh = y2 - y1
454
+ if bw <= 0 or bh <= 0:
455
+ continue
456
+ if bw < self.min_side or bh < self.min_side:
457
+ continue
458
+ area = bw * bh
459
+ if area < self.min_box_area:
460
+ continue
461
+ if area > 0.95 * image_area:
462
+ continue
463
+ ar = max(bw / max(bh, 1e-6), bh / max(bw, 1e-6))
464
+ if ar > self.max_aspect_ratio:
465
+ continue
466
+ keep.append(i)
467
+ if not keep:
468
+ return (
469
+ np.empty((0, 4), dtype=np.float32),
470
+ np.empty((0,), dtype=np.float32),
471
+ np.empty((0,), dtype=np.int32),
472
+ )
473
+ k = np.array(keep, dtype=np.intp)
474
+ return boxes[k], scores[k], cls_ids[k]
475
+
476
+ def _per_view_pipeline(
477
+ self,
478
+ boxes: np.ndarray,
479
+ scores: np.ndarray,
480
+ cls_ids: np.ndarray,
481
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
482
+ """Per-view post-processing pipeline: per-class NMS -> cap -> cross-class dedup."""
483
+ if len(boxes) > 1:
484
+ if self.use_soft_nms:
485
+ keep, new_scores = self._per_class_soft_nms(
486
+ boxes, scores, cls_ids,
487
+ self.soft_nms_sigma, self.soft_nms_score_thresh)
488
+ boxes, scores, cls_ids = boxes[keep], new_scores, cls_ids[keep]
489
+ else:
490
+ keep = self._per_class_hard_nms(boxes, scores, cls_ids, self.iou_thres)
491
+ boxes, scores, cls_ids = boxes[keep], scores[keep], cls_ids[keep]
492
+ if len(scores) > self.max_det:
493
+ top = np.argsort(-scores)[: self.max_det]
494
+ boxes, scores, cls_ids = boxes[top], scores[top], cls_ids[top]
495
+ if len(boxes) > 1:
496
+ boxes, scores, cls_ids = self._cross_class_dedup_op(
497
+ boxes, scores, cls_ids, self.cross_iou_thresh
498
+ )
499
+ return boxes, scores, cls_ids
500
+
501
+ @staticmethod
502
+ def _build_results(
503
+ boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray
504
+ ) -> list[BoundingBox]:
505
+ results: list[BoundingBox] = []
506
+ for box, conf, cls_id in zip(boxes, scores, cls_ids):
507
+ x1, y1, x2, y2 = box.tolist()
508
+ if x2 <= x1 or y2 <= y1:
509
+ continue
510
+ results.append(
511
+ BoundingBox(
512
+ x1=int(math.floor(x1)),
513
+ y1=int(math.floor(y1)),
514
+ x2=int(math.ceil(x2)),
515
+ y2=int(math.ceil(y2)),
516
+ cls_id=int(cls_id),
517
+ conf=float(conf),
518
+ )
519
+ )
520
+ return results
521
+
522
+ def _decode_final_dets(
523
+ self,
524
+ preds: np.ndarray,
525
+ ratio: float,
526
+ pad: tuple[float, float],
527
+ orig_size: tuple[int, int],
528
+ ) -> list[BoundingBox]:
529
+ """Final-detection output path: rows shaped [x1, y1, x2, y2, conf, cls_id]."""
530
+ if preds.ndim == 3 and preds.shape[0] == 1:
531
+ preds = preds[0]
532
+ if preds.ndim != 2 or preds.shape[1] < 6:
533
+ raise ValueError(f"Unexpected ONNX final-det output shape: {preds.shape}")
534
+
535
+ boxes = preds[:, :4].astype(np.float32)
536
+ scores = preds[:, 4].astype(np.float32)
537
+ cls_ids = preds[:, 5].astype(np.int32)
538
+ cls_ids = self.cls_remap[cls_ids]
539
+
540
+ keep = self._conf_filter_mask(scores, cls_ids)
541
+ boxes = boxes[keep]
542
+ scores = scores[keep]
543
+ cls_ids = cls_ids[keep]
544
+ if len(boxes) == 0:
545
+ return []
546
+
547
+ pad_w, pad_h = pad
548
+ boxes[:, [0, 2]] -= pad_w
549
+ boxes[:, [1, 3]] -= pad_h
550
+ boxes /= ratio
551
+ boxes = self._clip_boxes(boxes, orig_size)
552
+
553
+ boxes, scores, cls_ids = self._filter_sane_boxes(
554
+ boxes, scores, cls_ids, orig_size
555
+ )
556
+ if len(boxes) == 0:
557
+ return []
558
+
559
+ boxes, scores, cls_ids = self._per_view_pipeline(boxes, scores, cls_ids)
560
+ return self._build_results(boxes, scores, cls_ids)
561
+
562
+ def _decode_raw_yolo(
563
+ self,
564
+ preds: np.ndarray,
565
+ ratio: float,
566
+ pad: tuple[float, float],
567
+ orig_size: tuple[int, int],
568
+ ) -> list[BoundingBox]:
569
+ """Fallback raw-YOLO output path: per-anchor class logits."""
570
+ if preds.ndim != 3 or preds.shape[0] != 1:
571
+ raise ValueError(f"Unexpected raw ONNX output shape: {preds.shape}")
572
+ preds = preds[0]
573
+ if preds.shape[0] <= 16 and preds.shape[1] > preds.shape[0]:
574
+ preds = preds.T
575
+ if preds.ndim != 2 or preds.shape[1] < 5:
576
+ raise ValueError(f"Unexpected raw output shape: {preds.shape}")
577
+
578
+ boxes_xywh = preds[:, :4].astype(np.float32)
579
+ cls_part = preds[:, 4:].astype(np.float32)
580
+ if cls_part.shape[1] == 1:
581
+ scores = cls_part[:, 0]
582
+ cls_ids = np.zeros(len(scores), dtype=np.int32)
583
+ else:
584
+ cls_ids = np.argmax(cls_part, axis=1).astype(np.int32)
585
+ scores = cls_part[np.arange(len(cls_part)), cls_ids]
586
+ cls_ids = self.cls_remap[cls_ids]
587
+
588
+ keep = self._conf_filter_mask(scores, cls_ids)
589
+ boxes_xywh = boxes_xywh[keep]
590
+ scores = scores[keep]
591
+ cls_ids = cls_ids[keep]
592
+ if len(boxes_xywh) == 0:
593
+ return []
594
+ boxes = self._xywh_to_xyxy(boxes_xywh)
595
+
596
+ pad_w, pad_h = pad
597
+ boxes[:, [0, 2]] -= pad_w
598
+ boxes[:, [1, 3]] -= pad_h
599
+ boxes /= ratio
600
+ boxes = self._clip_boxes(boxes, orig_size)
601
+
602
+ boxes, scores, cls_ids = self._filter_sane_boxes(
603
+ boxes, scores, cls_ids, orig_size
604
+ )
605
+ if len(boxes) == 0:
606
+ return []
607
+
608
+ boxes, scores, cls_ids = self._per_view_pipeline(boxes, scores, cls_ids)
609
+ return self._build_results(boxes, scores, cls_ids)
610
+
611
+ def _postprocess(
612
+ self,
613
+ output: np.ndarray,
614
+ ratio: float,
615
+ pad: tuple[float, float],
616
+ orig_size: tuple[int, int],
617
+ ) -> list[BoundingBox]:
618
+ if output.ndim == 2 and output.shape[1] >= 6:
619
+ return self._decode_final_dets(output, ratio, pad, orig_size)
620
+ if output.ndim == 3 and output.shape[0] == 1 and output.shape[2] == 6:
621
+ return self._decode_final_dets(output, ratio, pad, orig_size)
622
+ return self._decode_raw_yolo(output, ratio, pad, orig_size)
623
+
624
+ def _predict_single(self, image: np.ndarray) -> list[BoundingBox]:
625
+ if image is None:
626
+ raise ValueError("Input image is None")
627
+ if not isinstance(image, np.ndarray):
628
+ raise TypeError(f"Input is not numpy array: {type(image)}")
629
+ if image.ndim != 3:
630
+ raise ValueError(f"Expected HWC image, got shape={image.shape}")
631
+ if image.shape[0] <= 0 or image.shape[1] <= 0:
632
+ raise ValueError(f"Invalid image shape={image.shape}")
633
+ if image.shape[2] != 3:
634
+ raise ValueError(f"Expected 3 channels, got shape={image.shape}")
635
+ if image.dtype != np.uint8:
636
+ image = image.astype(np.uint8)
637
+
638
+ input_tensor, ratio, pad, orig_size = self._preprocess(image)
639
+ expected = (1, 3, self.input_height, self.input_width)
640
+ if input_tensor.shape != expected:
641
+ raise ValueError(
642
+ f"Bad input tensor shape={input_tensor.shape}, expected={expected}"
643
+ )
644
+
645
+ outputs = self.session.run(self.output_names, {self.input_name: input_tensor})
646
+ return self._postprocess(outputs[0], ratio, pad, orig_size)
647
+
648
+ def _predict_tta(self, image: np.ndarray) -> list[BoundingBox]:
649
+ """Horizontal-flip TTA.
650
+
651
+ Strategy:
652
+ 1. Predict on original and on flipped image.
653
+ 2. Map flipped boxes back to original coordinates.
654
+ 3. Per-class hard NMS on the union.
655
+ 4. For each kept box, compute the max same-class score across the
656
+ FULL union (not just the post-NMS subset) -- this lets a high-
657
+ confidence flipped detection raise a borderline original one.
658
+ 5. Cross-class dedup to suppress same-physical-object multi-class.
659
+ """
660
+ boxes_orig = self._predict_single(image)
661
+ flipped = cv2.flip(image, 1)
662
+ boxes_flip = self._predict_single(flipped)
663
+ w = image.shape[1]
664
+ boxes_flip = [
665
+ BoundingBox(
666
+ x1=w - b.x2, y1=b.y1, x2=w - b.x1, y2=b.y2,
667
+ cls_id=b.cls_id, conf=b.conf,
668
+ )
669
+ for b in boxes_flip
670
+ ]
671
+ all_boxes = boxes_orig + boxes_flip
672
+ if not all_boxes:
673
+ return []
674
+
675
+ coords = np.array(
676
+ [[b.x1, b.y1, b.x2, b.y2] for b in all_boxes], dtype=np.float32
677
+ )
678
+ scores = np.array([b.conf for b in all_boxes], dtype=np.float32)
679
+ cls_ids = np.array([b.cls_id for b in all_boxes], dtype=np.int32)
680
+
681
+ hard_keep = self._per_class_hard_nms(coords, scores, cls_ids, self.iou_thres)
682
+ if len(hard_keep) == 0:
683
+ return []
684
+ if len(hard_keep) > self.max_det:
685
+ top = np.argsort(-scores[hard_keep])[: self.max_det]
686
+ hard_keep = hard_keep[top]
687
+
688
+ boosted = self._max_score_per_cluster(
689
+ coords[hard_keep], cls_ids[hard_keep],
690
+ coords, scores, cls_ids, self.iou_thres,
691
+ )
692
+
693
+ kept_coords = coords[hard_keep]
694
+ kept_cls = cls_ids[hard_keep]
695
+ if len(kept_coords) > 1:
696
+ kept_coords, boosted, kept_cls = self._cross_class_dedup_op(
697
+ kept_coords, boosted, kept_cls, self.cross_iou_thresh
698
+ )
699
+
700
+ return [
701
+ BoundingBox(
702
+ x1=int(math.floor(kept_coords[j, 0])),
703
+ y1=int(math.floor(kept_coords[j, 1])),
704
+ x2=int(math.ceil(kept_coords[j, 2])),
705
+ y2=int(math.ceil(kept_coords[j, 3])),
706
+ cls_id=int(kept_cls[j]),
707
+ conf=float(boosted[j]),
708
+ )
709
+ for j in range(len(kept_coords))
710
+ ]
711
+
712
+ def _predict_tiles(self, image: np.ndarray) -> list[BoundingBox]:
713
+ """Tile-based TTA for high-resolution images.
714
+
715
+ Splits the source image into two overlapping horizontal tiles, runs
716
+ single-pass inference on each at native scale, and translates boxes
717
+ back to the global frame. Useful when source width >> model input
718
+ width because letterboxing otherwise discards effective resolution
719
+ that small / distant signs depend on.
720
+
721
+ Returns an empty list if the image isn't wide enough to benefit; the
722
+ caller falls back to the regular pipeline in that case.
723
+ """
724
+ h, w = image.shape[:2]
725
+ if w < int(self.input_width * self.tile_trigger_ratio):
726
+ return []
727
+
728
+ overlap = int(w * self.tile_overlap_ratio)
729
+ mid = w // 2
730
+ x_left_end = min(w, mid + overlap // 2)
731
+ x_right_start = max(0, mid - overlap // 2)
732
+
733
+ left = image[:, :x_left_end]
734
+ right = image[:, x_right_start:]
735
+
736
+ boxes_left = self._predict_single(left)
737
+ boxes_right = self._predict_single(right)
738
+
739
+ shifted_right = [
740
+ BoundingBox(
741
+ x1=b.x1 + x_right_start,
742
+ y1=b.y1,
743
+ x2=b.x2 + x_right_start,
744
+ y2=b.y2,
745
+ cls_id=b.cls_id,
746
+ conf=b.conf,
747
+ )
748
+ for b in boxes_right
749
+ ]
750
+ return boxes_left + shifted_right
751
+
752
+ def _merge_views(
753
+ self,
754
+ view_boxes: list[list[BoundingBox]],
755
+ image_size: tuple[int, int],
756
+ ) -> list[BoundingBox]:
757
+ """Merge boxes from multiple views (single / hflip / tiles).
758
+
759
+ Same logic as `_predict_tta`'s tail: per-class hard NMS to dedupe,
760
+ then for each kept box take the max same-class score across the full
761
+ candidate union — a high-confidence detection in any view boosts
762
+ borderline matches in others.
763
+ """
764
+ all_boxes: list[BoundingBox] = []
765
+ for vb in view_boxes:
766
+ all_boxes.extend(vb)
767
+ if not all_boxes:
768
+ return []
769
+
770
+ coords = np.array(
771
+ [[b.x1, b.y1, b.x2, b.y2] for b in all_boxes], dtype=np.float32
772
+ )
773
+ scores = np.array([b.conf for b in all_boxes], dtype=np.float32)
774
+ cls_ids = np.array([b.cls_id for b in all_boxes], dtype=np.int32)
775
+
776
+ coords = self._clip_boxes(coords, image_size)
777
+
778
+ hard_keep = self._per_class_hard_nms(coords, scores, cls_ids, self.iou_thres)
779
+ if len(hard_keep) == 0:
780
+ return []
781
+ if len(hard_keep) > self.max_det:
782
+ top = np.argsort(-scores[hard_keep])[: self.max_det]
783
+ hard_keep = hard_keep[top]
784
+
785
+ boosted = self._max_score_per_cluster(
786
+ coords[hard_keep], cls_ids[hard_keep],
787
+ coords, scores, cls_ids, self.iou_thres,
788
+ )
789
+
790
+ kept_coords = coords[hard_keep]
791
+ kept_cls = cls_ids[hard_keep]
792
+ if len(kept_coords) > 1:
793
+ kept_coords, boosted, kept_cls = self._cross_class_dedup_op(
794
+ kept_coords, boosted, kept_cls, self.cross_iou_thresh
795
+ )
796
+
797
+ return [
798
+ BoundingBox(
799
+ x1=int(math.floor(kept_coords[j, 0])),
800
+ y1=int(math.floor(kept_coords[j, 1])),
801
+ x2=int(math.ceil(kept_coords[j, 2])),
802
+ y2=int(math.ceil(kept_coords[j, 3])),
803
+ cls_id=int(kept_cls[j]),
804
+ conf=float(boosted[j]),
805
+ )
806
+ for j in range(len(kept_coords))
807
+ ]
808
+
809
+ def _predict_full(self, image: np.ndarray) -> list[BoundingBox]:
810
+ """Top-level per-frame prediction with all enabled augmentations.
811
+
812
+ - `use_tta=True`: original + horizontal flip
813
+ - `use_tile_tta=True` AND image wide enough: two overlapping tiles
814
+ All views are merged via per-class NMS + cluster-max score boost.
815
+ """
816
+ if not self.use_tta and not self.use_tile_tta:
817
+ return self._predict_single(image)
818
+
819
+ views: list[list[BoundingBox]] = []
820
+ if self.use_tta:
821
+ views.append(self._predict_single(image))
822
+ flipped = cv2.flip(image, 1)
823
+ w = image.shape[1]
824
+ flipped_dets = self._predict_single(flipped)
825
+ views.append([
826
+ BoundingBox(
827
+ x1=w - b.x2, y1=b.y1, x2=w - b.x1, y2=b.y2,
828
+ cls_id=b.cls_id, conf=b.conf,
829
+ )
830
+ for b in flipped_dets
831
+ ])
832
+ else:
833
+ views.append(self._predict_single(image))
834
+
835
+ if self.use_tile_tta:
836
+ tile_boxes = self._predict_tiles(image)
837
+ if tile_boxes:
838
+ views.append(tile_boxes)
839
+
840
+ h, w = image.shape[:2]
841
+ return self._merge_views(views, (w, h))
842
+
843
+ def predict_batch(
844
+ self,
845
+ batch_images: list[ndarray],
846
+ offset: int,
847
+ n_keypoints: int,
848
+ ) -> list[TVFrameResult]:
849
+ results: list[TVFrameResult] = []
850
+ for frame_number_in_batch, image in enumerate(batch_images):
851
+ try:
852
+ boxes = self._predict_full(image)
853
+ except Exception as e:
854
+ print(
855
+ f"⚠️ Inference failed for frame "
856
+ f"{offset + frame_number_in_batch}: {e}"
857
+ )
858
+ boxes = []
859
+ results.append(
860
+ TVFrameResult(
861
+ frame_id=offset + frame_number_in_batch,
862
+ boxes=boxes,
863
+ keypoints=[(0, 0) for _ in range(max(0, int(n_keypoints)))],
864
+ )
865
+ )
866
+ return results
weights.onnx ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:2c530c38045a23d27336666bf1adde01ec90417276a4ab59207d36e2409e5f3f
3
+ size 19383133