nexu02 commited on
Commit
86cb5ab
·
0 Parent(s):

chore: squash history (remove leaked training artifacts)

Browse files
Files changed (5) hide show
  1. .gitattributes +35 -0
  2. chute_config.yml +21 -0
  3. miner.py +642 -0
  4. readme.md +1 -0
  5. weights.onnx +3 -0
.gitattributes ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ *.7z filter=lfs diff=lfs merge=lfs -text
2
+ *.arrow filter=lfs diff=lfs merge=lfs -text
3
+ *.bin filter=lfs diff=lfs merge=lfs -text
4
+ *.bz2 filter=lfs diff=lfs merge=lfs -text
5
+ *.ckpt filter=lfs diff=lfs merge=lfs -text
6
+ *.ftz filter=lfs diff=lfs merge=lfs -text
7
+ *.gz filter=lfs diff=lfs merge=lfs -text
8
+ *.h5 filter=lfs diff=lfs merge=lfs -text
9
+ *.joblib filter=lfs diff=lfs merge=lfs -text
10
+ *.lfs.* filter=lfs diff=lfs merge=lfs -text
11
+ *.mlmodel filter=lfs diff=lfs merge=lfs -text
12
+ *.model filter=lfs diff=lfs merge=lfs -text
13
+ *.msgpack filter=lfs diff=lfs merge=lfs -text
14
+ *.npy filter=lfs diff=lfs merge=lfs -text
15
+ *.npz filter=lfs diff=lfs merge=lfs -text
16
+ *.onnx filter=lfs diff=lfs merge=lfs -text
17
+ *.ot filter=lfs diff=lfs merge=lfs -text
18
+ *.parquet filter=lfs diff=lfs merge=lfs -text
19
+ *.pb filter=lfs diff=lfs merge=lfs -text
20
+ *.pickle filter=lfs diff=lfs merge=lfs -text
21
+ *.pkl filter=lfs diff=lfs merge=lfs -text
22
+ *.pt filter=lfs diff=lfs merge=lfs -text
23
+ *.pth filter=lfs diff=lfs merge=lfs -text
24
+ *.rar filter=lfs diff=lfs merge=lfs -text
25
+ *.safetensors filter=lfs diff=lfs merge=lfs -text
26
+ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
27
+ *.tar.* filter=lfs diff=lfs merge=lfs -text
28
+ *.tar filter=lfs diff=lfs merge=lfs -text
29
+ *.tflite filter=lfs diff=lfs merge=lfs -text
30
+ *.tgz filter=lfs diff=lfs merge=lfs -text
31
+ *.wasm filter=lfs diff=lfs merge=lfs -text
32
+ *.xz filter=lfs diff=lfs merge=lfs -text
33
+ *.zip filter=lfs diff=lfs merge=lfs -text
34
+ *.zst filter=lfs diff=lfs merge=lfs -text
35
+ *tfevents* filter=lfs diff=lfs merge=lfs -text
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,642 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ # === Auto-generated postproc config — overridden by build_submission.py.
12
+ # Keys / value types must match what build_submission expects so the
13
+ # search outcomes from utils/hp_search.py transfer 1:1 to deploy.
14
+ CONFIG = {
15
+ 'class_names' : ["road sign"],
16
+ 'model_class_order' : ["road sign"],
17
+ 'input_size' : 1280,
18
+ 'conf_thres' : 0.45,
19
+ 'conf_thres_per_class': {"road sign": 0.45},
20
+ 'iou_thres' : 0.4,
21
+ 'nms_mode' : "hard_global",
22
+ 'soft_sigma' : 0.4,
23
+ 'cross_class_dedup' : True,
24
+ 'cross_iou_thresh' : 0.8,
25
+ 'sanity_filter' : False,
26
+ 'min_side' : 0.0,
27
+ 'min_box_area' : 0.0,
28
+ 'max_aspect_ratio' : 6.0,
29
+ 'max_det' : 300,
30
+ 'use_tta' : True,
31
+ 'bonus_per_class' : {"road sign": 0.0},
32
+ }
33
+
34
+
35
+ class BoundingBox(BaseModel):
36
+ x1: int
37
+ y1: int
38
+ x2: int
39
+ y2: int
40
+ cls_id: int
41
+ conf: float
42
+
43
+
44
+ class TVFrameResult(BaseModel):
45
+ frame_id: int
46
+ boxes: list[BoundingBox]
47
+ keypoints: list[tuple[int, int]]
48
+
49
+
50
+ class Miner:
51
+ """ONNX Runtime miner. Postproc is driven by the CONFIG block above."""
52
+
53
+ def __init__(self, path_hf_repo: Path) -> None:
54
+ cfg = CONFIG
55
+ model_path = path_hf_repo / "weights.onnx"
56
+ self.class_names = list(cfg["class_names"])
57
+ model_class_order = list(cfg["model_class_order"])
58
+ self.cls_remap = np.array(
59
+ [self.class_names.index(n) for n in model_class_order], dtype=np.int32
60
+ )
61
+ print("ORT version:", ort.__version__)
62
+
63
+ try:
64
+ ort.preload_dlls()
65
+ print("preload_dlls success")
66
+ except Exception as e:
67
+ print(f"preload_dlls failed: {e}")
68
+
69
+ print("ORT available providers BEFORE session:", ort.get_available_providers())
70
+
71
+ sess_options = ort.SessionOptions()
72
+ sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
73
+
74
+ try:
75
+ self.session = ort.InferenceSession(
76
+ str(model_path),
77
+ sess_options=sess_options,
78
+ providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
79
+ )
80
+ print("Created ORT session with preferred CUDA provider list")
81
+ except Exception as e:
82
+ print(f"CUDA session creation failed, falling back to CPU: {e}")
83
+ self.session = ort.InferenceSession(
84
+ str(model_path),
85
+ sess_options=sess_options,
86
+ providers=["CPUExecutionProvider"],
87
+ )
88
+
89
+ print("ORT session providers:", self.session.get_providers())
90
+
91
+ for inp in self.session.get_inputs():
92
+ print("INPUT:", inp.name, inp.shape, inp.type)
93
+ for out in self.session.get_outputs():
94
+ print("OUTPUT:", out.name, out.shape, out.type)
95
+
96
+ self.input_name = self.session.get_inputs()[0].name
97
+ self.output_names = [output.name for output in self.session.get_outputs()]
98
+ self.input_shape = self.session.get_inputs()[0].shape
99
+
100
+ default_size = int(cfg["input_size"])
101
+ self.input_height = self._safe_dim(self.input_shape[2], default=default_size)
102
+ self.input_width = self._safe_dim(self.input_shape[3], default=default_size)
103
+
104
+ self.conf_thres = float(cfg["conf_thres"])
105
+ per_class_cfg = cfg.get("conf_thres_per_class") or {}
106
+ self.conf_thres_per_class = {
107
+ str(k): float(v) for k, v in per_class_cfg.items()
108
+ }
109
+ # Indexed by remapped cls_id (i.e. position in self.class_names).
110
+ # Falls back to the global conf_thres for any class not in the dict.
111
+ self._conf_thres_array = np.array(
112
+ [self.conf_thres_per_class.get(name, self.conf_thres)
113
+ for name in self.class_names],
114
+ dtype=np.float32,
115
+ )
116
+ self.iou_thres = float(cfg["iou_thres"])
117
+ self.nms_mode = str(cfg["nms_mode"])
118
+ self.soft_sigma = float(cfg["soft_sigma"])
119
+ self.cross_class_dedup = bool(cfg["cross_class_dedup"])
120
+ self.cross_iou_thresh = float(cfg["cross_iou_thresh"])
121
+ self.sanity_filter = bool(cfg["sanity_filter"])
122
+ self.min_side = float(cfg["min_side"])
123
+ self.min_box_area = float(cfg["min_box_area"])
124
+ self.max_aspect_ratio = float(cfg["max_aspect_ratio"])
125
+ self.max_det = int(cfg["max_det"])
126
+ self.use_tta = bool(cfg["use_tta"])
127
+ bonus_cfg = cfg.get("bonus_per_class") or {}
128
+ self.bonus_per_class = {str(k): float(v) for k, v in bonus_cfg.items()}
129
+ # Indexed by remapped cls_id (position in self.class_names).
130
+ self._bonus_array = np.array(
131
+ [self.bonus_per_class.get(name, 0.0) for name in self.class_names],
132
+ dtype=np.float32,
133
+ )
134
+
135
+ print(f"ONNX model loaded from: {model_path}")
136
+ print(f"ONNX providers: {self.session.get_providers()}")
137
+ print(f"ONNX input: name={self.input_name}, shape={self.input_shape}")
138
+ print(f"postproc: conf={self.conf_thres} iou={self.iou_thres} "
139
+ f"nms={self.nms_mode} dedup={self.cross_class_dedup} "
140
+ f"sanity={self.sanity_filter} tta={self.use_tta} "
141
+ f"max_det={self.max_det}")
142
+ print("per-class conf: " + ", ".join(
143
+ f"{n}={t:.3f}" for n, t in zip(self.class_names,
144
+ self._conf_thres_array.tolist())))
145
+
146
+ def __repr__(self) -> str:
147
+ return (
148
+ f"ONNXRuntime(session={type(self.session).__name__}, "
149
+ f"providers={self.session.get_providers()})"
150
+ )
151
+
152
+ @staticmethod
153
+ def _safe_dim(value, default: int) -> int:
154
+ return value if isinstance(value, int) and value > 0 else default
155
+
156
+ def _letterbox(self, image: ndarray, new_shape: tuple[int, int],
157
+ color=(114, 114, 114)
158
+ ) -> tuple[ndarray, float, tuple[float, float]]:
159
+ h, w = image.shape[:2]
160
+ new_w, new_h = new_shape
161
+ ratio = min(new_w / w, new_h / h)
162
+ resized_w = int(round(w * ratio))
163
+ resized_h = int(round(h * ratio))
164
+ if (resized_w, resized_h) != (w, h):
165
+ interp = cv2.INTER_CUBIC if ratio > 1.0 else cv2.INTER_LINEAR
166
+ image = cv2.resize(image, (resized_w, resized_h), interpolation=interp)
167
+ dw = (new_w - resized_w) / 2.0
168
+ dh = (new_h - resized_h) / 2.0
169
+ left = int(round(dw - 0.1))
170
+ right = int(round(dw + 0.1))
171
+ top = int(round(dh - 0.1))
172
+ bottom = int(round(dh + 0.1))
173
+ padded = cv2.copyMakeBorder(image, top, bottom, left, right,
174
+ borderType=cv2.BORDER_CONSTANT, value=color)
175
+ return padded, ratio, (dw, dh)
176
+
177
+ def _preprocess(self, image: ndarray
178
+ ) -> tuple[np.ndarray, float, tuple[float, float],
179
+ tuple[int, int]]:
180
+ orig_h, orig_w = image.shape[:2]
181
+ img, ratio, pad = self._letterbox(image, (self.input_width, self.input_height))
182
+ img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
183
+ img = img.astype(np.float32) / 255.0
184
+ img = np.transpose(img, (2, 0, 1))[None, ...]
185
+ img = np.ascontiguousarray(img, dtype=np.float32)
186
+ return img, ratio, pad, (orig_w, orig_h)
187
+
188
+ @staticmethod
189
+ def _clip_boxes(boxes: np.ndarray, image_size: tuple[int, int]) -> np.ndarray:
190
+ w, h = image_size
191
+ boxes[:, 0] = np.clip(boxes[:, 0], 0, w - 1)
192
+ boxes[:, 1] = np.clip(boxes[:, 1], 0, h - 1)
193
+ boxes[:, 2] = np.clip(boxes[:, 2], 0, w - 1)
194
+ boxes[:, 3] = np.clip(boxes[:, 3], 0, h - 1)
195
+ return boxes
196
+
197
+ @staticmethod
198
+ def _xywh_to_xyxy(boxes: np.ndarray) -> np.ndarray:
199
+ out = np.empty_like(boxes)
200
+ out[:, 0] = boxes[:, 0] - boxes[:, 2] / 2.0
201
+ out[:, 1] = boxes[:, 1] - boxes[:, 3] / 2.0
202
+ out[:, 2] = boxes[:, 0] + boxes[:, 2] / 2.0
203
+ out[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.0
204
+ return out
205
+
206
+ @staticmethod
207
+ def _hard_nms(boxes: np.ndarray, scores: np.ndarray,
208
+ iou_thresh: float) -> np.ndarray:
209
+ n = len(boxes)
210
+ if n == 0:
211
+ return np.array([], dtype=np.intp)
212
+ order = np.argsort(-scores)
213
+ keep: list[int] = []
214
+ while len(order) > 0:
215
+ i = int(order[0])
216
+ keep.append(i)
217
+ if len(order) == 1:
218
+ break
219
+ rest = order[1:]
220
+ xx1 = np.maximum(boxes[i, 0], boxes[rest, 0])
221
+ yy1 = np.maximum(boxes[i, 1], boxes[rest, 1])
222
+ xx2 = np.minimum(boxes[i, 2], boxes[rest, 2])
223
+ yy2 = np.minimum(boxes[i, 3], boxes[rest, 3])
224
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
225
+ a_i = (max(0.0, boxes[i, 2] - boxes[i, 0]) *
226
+ max(0.0, boxes[i, 3] - boxes[i, 1]))
227
+ a_r = (np.maximum(0.0, boxes[rest, 2] - boxes[rest, 0]) *
228
+ np.maximum(0.0, boxes[rest, 3] - boxes[rest, 1]))
229
+ iou = inter / (a_i + a_r - inter + 1e-7)
230
+ order = rest[iou <= iou_thresh]
231
+ return np.array(keep, dtype=np.intp)
232
+
233
+ @staticmethod
234
+ def _soft_nms(boxes: np.ndarray, scores: np.ndarray,
235
+ sigma: float = 0.5, score_thresh: float = 0.001
236
+ ) -> tuple[np.ndarray, np.ndarray]:
237
+ n = len(boxes)
238
+ if n == 0:
239
+ return np.array([], dtype=np.intp), np.array([], dtype=np.float32)
240
+ boxes = boxes.astype(np.float32, copy=True)
241
+ scores = scores.astype(np.float32, copy=True)
242
+ order = np.arange(n)
243
+ for i in range(n):
244
+ max_pos = i + int(np.argmax(scores[i:]))
245
+ boxes[[i, max_pos]] = boxes[[max_pos, i]]
246
+ scores[[i, max_pos]] = scores[[max_pos, i]]
247
+ order[[i, max_pos]] = order[[max_pos, i]]
248
+ if i + 1 >= n:
249
+ break
250
+ xx1 = np.maximum(boxes[i, 0], boxes[i + 1:, 0])
251
+ yy1 = np.maximum(boxes[i, 1], boxes[i + 1:, 1])
252
+ xx2 = np.minimum(boxes[i, 2], boxes[i + 1:, 2])
253
+ yy2 = np.minimum(boxes[i, 3], boxes[i + 1:, 3])
254
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
255
+ a_i = max(0.0, float((boxes[i, 2] - boxes[i, 0]) *
256
+ (boxes[i, 3] - boxes[i, 1])))
257
+ a_j = (np.maximum(0.0, boxes[i + 1:, 2] - boxes[i + 1:, 0]) *
258
+ np.maximum(0.0, boxes[i + 1:, 3] - boxes[i + 1:, 1]))
259
+ iou = inter / (a_i + a_j - inter + 1e-7)
260
+ scores[i + 1:] *= np.exp(-(iou ** 2) / sigma)
261
+ mask = scores > score_thresh
262
+ return order[mask], scores[mask]
263
+
264
+ def _per_class_hard_nms(self, boxes: np.ndarray, scores: np.ndarray,
265
+ cls_ids: np.ndarray, iou_thresh: float
266
+ ) -> np.ndarray:
267
+ if len(boxes) == 0:
268
+ return np.array([], dtype=np.intp)
269
+ all_keep: list[int] = []
270
+ for c in np.unique(cls_ids):
271
+ mask = cls_ids == c
272
+ indices = np.where(mask)[0]
273
+ keep = self._hard_nms(boxes[mask], scores[mask], iou_thresh)
274
+ all_keep.extend(indices[keep].tolist())
275
+ all_keep.sort()
276
+ return np.array(all_keep, dtype=np.intp)
277
+
278
+ def _per_class_soft_nms(self, boxes: np.ndarray, scores: np.ndarray,
279
+ cls_ids: np.ndarray, sigma: float
280
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
281
+ if len(boxes) == 0:
282
+ return boxes, scores, cls_ids
283
+ out_b: list = []
284
+ out_s: list = []
285
+ out_c: list = []
286
+ for c in np.unique(cls_ids):
287
+ mask = cls_ids == c
288
+ sub_b = boxes[mask]
289
+ sub_s = scores[mask]
290
+ sub_c = cls_ids[mask]
291
+ idx, decayed = self._soft_nms(sub_b, sub_s, sigma)
292
+ if len(idx) == 0:
293
+ continue
294
+ out_b.append(sub_b[idx])
295
+ out_s.append(decayed)
296
+ out_c.append(sub_c[idx])
297
+ if not out_b:
298
+ return (np.empty((0, 4), dtype=np.float32),
299
+ np.empty((0,), dtype=np.float32),
300
+ np.empty((0,), dtype=cls_ids.dtype))
301
+ return (np.concatenate(out_b, axis=0),
302
+ np.concatenate(out_s, axis=0),
303
+ np.concatenate(out_c, axis=0))
304
+
305
+ def _apply_nms(self, boxes: np.ndarray, scores: np.ndarray,
306
+ cls_ids: np.ndarray, iou_thresh: float
307
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
308
+ if self.nms_mode == "hard_global":
309
+ keep = self._hard_nms(boxes, scores, iou_thresh)
310
+ return boxes[keep], scores[keep], cls_ids[keep]
311
+ if self.nms_mode == "hard_per_class":
312
+ keep = self._per_class_hard_nms(boxes, scores, cls_ids, iou_thresh)
313
+ return boxes[keep], scores[keep], cls_ids[keep]
314
+ if self.nms_mode == "soft_per_class":
315
+ return self._per_class_soft_nms(boxes, scores, cls_ids, self.soft_sigma)
316
+ raise ValueError(f"unknown nms_mode={self.nms_mode!r}")
317
+
318
+ def _cross_class_dedup_op(self, boxes: np.ndarray, scores: np.ndarray,
319
+ cls_ids: np.ndarray, iou_thresh: float
320
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
321
+ n = len(boxes)
322
+ if n <= 1:
323
+ return boxes, scores, cls_ids
324
+ boxes = np.asarray(boxes, dtype=np.float32)
325
+ scores = np.asarray(scores, dtype=np.float32)
326
+ cls_ids = np.asarray(cls_ids, dtype=np.int32)
327
+ areas = (np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) *
328
+ np.maximum(0.0, boxes[:, 3] - boxes[:, 1]))
329
+ margins = scores - self._conf_thres_array[cls_ids]
330
+ order = np.lexsort((-areas, -margins))
331
+ suppressed = np.zeros(n, dtype=bool)
332
+ keep: list[int] = []
333
+ for i in order:
334
+ if suppressed[i]:
335
+ continue
336
+ keep.append(int(i))
337
+ bi = boxes[i]
338
+ xx1 = np.maximum(bi[0], boxes[:, 0])
339
+ yy1 = np.maximum(bi[1], boxes[:, 1])
340
+ xx2 = np.minimum(bi[2], boxes[:, 2])
341
+ yy2 = np.minimum(bi[3], boxes[:, 3])
342
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
343
+ a_i = max(1e-7, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
344
+ iou = inter / (a_i + areas - inter + 1e-7)
345
+ dup = iou > iou_thresh
346
+ dup[i] = False
347
+ suppressed |= dup
348
+ keep_idx = np.array(keep, dtype=np.intp)
349
+ return boxes[keep_idx], scores[keep_idx], cls_ids[keep_idx]
350
+
351
+ def _filter_sane_boxes_op(self, boxes: np.ndarray, scores: np.ndarray,
352
+ cls_ids: np.ndarray, orig_size: tuple[int, int]
353
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
354
+ if len(boxes) == 0:
355
+ return boxes, scores, cls_ids
356
+ orig_w, orig_h = orig_size
357
+ image_area = float(orig_w * orig_h)
358
+ bw = np.maximum(0.0, boxes[:, 2] - boxes[:, 0])
359
+ bh = np.maximum(0.0, boxes[:, 3] - boxes[:, 1])
360
+ area = bw * bh
361
+ ar = np.where(
362
+ (bw > 0) & (bh > 0),
363
+ np.maximum(bw / np.maximum(bh, 1e-6), bh / np.maximum(bw, 1e-6)),
364
+ np.inf,
365
+ )
366
+ keep = (
367
+ (bw >= self.min_side) & (bh >= self.min_side) &
368
+ (area >= self.min_box_area) &
369
+ (area <= 0.95 * image_area) &
370
+ (ar <= self.max_aspect_ratio)
371
+ )
372
+ return boxes[keep], scores[keep], cls_ids[keep]
373
+
374
+ @staticmethod
375
+ def _max_score_per_cluster(post_boxes: np.ndarray,
376
+ post_cls: np.ndarray,
377
+ full_boxes: np.ndarray,
378
+ full_scores: np.ndarray,
379
+ full_cls: np.ndarray,
380
+ iou_thresh: float) -> np.ndarray:
381
+ n = len(post_boxes)
382
+ if n == 0:
383
+ return np.empty(0, dtype=np.float32)
384
+ full_areas = (np.maximum(0.0, full_boxes[:, 2] - full_boxes[:, 0]) *
385
+ np.maximum(0.0, full_boxes[:, 3] - full_boxes[:, 1]))
386
+ out = np.empty(n, dtype=np.float32)
387
+ for i in range(n):
388
+ bi = post_boxes[i]
389
+ xx1 = np.maximum(bi[0], full_boxes[:, 0])
390
+ yy1 = np.maximum(bi[1], full_boxes[:, 1])
391
+ xx2 = np.minimum(bi[2], full_boxes[:, 2])
392
+ yy2 = np.minimum(bi[3], full_boxes[:, 3])
393
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
394
+ a_i = max(0.0, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
395
+ iou = inter / (a_i + full_areas - inter + 1e-7)
396
+ cluster = (iou >= iou_thresh) & (full_cls == post_cls[i])
397
+ out[i] = float(np.max(full_scores[cluster])) if np.any(cluster) else 0.0
398
+ return out
399
+
400
+ def _conf_filter_mask(self, scores: np.ndarray,
401
+ cls_ids: np.ndarray) -> np.ndarray:
402
+ """Boolean keep-mask: score >= per-class threshold, with a per-class
403
+ rescue — if a class has zero boxes passing, admit its top-1 candidate
404
+ when its score >= (per-class threshold - per-class bonus). A class
405
+ with bonus=0 is unaffected (legacy behavior)."""
406
+ if len(scores) == 0:
407
+ return np.zeros(0, dtype=bool)
408
+ thr = self._conf_thres_array[cls_ids]
409
+ keep = scores >= thr
410
+ for c in np.unique(cls_ids):
411
+ b = float(self._bonus_array[c])
412
+ if b <= 0.0:
413
+ continue
414
+ cm = cls_ids == c
415
+ if keep[cm].any():
416
+ continue
417
+ idx = np.where(cm)[0]
418
+ top = int(idx[int(np.argmax(scores[idx]))])
419
+ if scores[top] >= self._conf_thres_array[c] - b:
420
+ keep[top] = True
421
+ return keep
422
+
423
+ def _per_view_pipeline(self, boxes: np.ndarray, scores: np.ndarray,
424
+ cls_ids: np.ndarray, orig_size: tuple[int, int]
425
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
426
+ """Apply the per-view postproc pipeline (post conf-filter):
427
+ sanity (if on) → NMS → max_det → cross_dedup (if on)."""
428
+ if self.sanity_filter:
429
+ boxes, scores, cls_ids = self._filter_sane_boxes_op(
430
+ boxes, scores, cls_ids, orig_size
431
+ )
432
+ if len(boxes) == 0:
433
+ return boxes, scores, cls_ids
434
+ if len(boxes) > 1:
435
+ boxes, scores, cls_ids = self._apply_nms(
436
+ boxes, scores, cls_ids, self.iou_thres
437
+ )
438
+ if len(scores) > self.max_det:
439
+ top = np.argsort(-scores)[: self.max_det]
440
+ boxes, scores, cls_ids = boxes[top], scores[top], cls_ids[top]
441
+ if self.cross_class_dedup and len(boxes) > 1:
442
+ boxes, scores, cls_ids = self._cross_class_dedup_op(
443
+ boxes, scores, cls_ids, self.cross_iou_thresh
444
+ )
445
+ return boxes, scores, cls_ids
446
+
447
+ def _decode_final_dets(self, preds: np.ndarray, ratio: float,
448
+ pad: tuple[float, float],
449
+ orig_size: tuple[int, int]) -> list[BoundingBox]:
450
+ if preds.ndim == 3 and preds.shape[0] == 1:
451
+ preds = preds[0]
452
+ if preds.ndim != 2 or preds.shape[1] < 6:
453
+ raise ValueError(f"Unexpected ONNX final-det output shape: {preds.shape}")
454
+
455
+ boxes = preds[:, :4].astype(np.float32)
456
+ scores = preds[:, 4].astype(np.float32)
457
+ cls_ids = preds[:, 5].astype(np.int32)
458
+ cls_ids = self.cls_remap[cls_ids]
459
+
460
+ keep = self._conf_filter_mask(scores, cls_ids)
461
+ boxes = boxes[keep]
462
+ scores = scores[keep]
463
+ cls_ids = cls_ids[keep]
464
+ if len(boxes) == 0:
465
+ return []
466
+
467
+ pad_w, pad_h = pad
468
+ boxes[:, [0, 2]] -= pad_w
469
+ boxes[:, [1, 3]] -= pad_h
470
+ boxes /= ratio
471
+ boxes = self._clip_boxes(boxes, orig_size)
472
+
473
+ boxes, scores, cls_ids = self._per_view_pipeline(
474
+ boxes, scores, cls_ids, orig_size
475
+ )
476
+ return self._build_results(boxes, scores, cls_ids)
477
+
478
+ def _decode_raw_yolo(self, preds: np.ndarray, ratio: float,
479
+ pad: tuple[float, float],
480
+ orig_size: tuple[int, int]) -> list[BoundingBox]:
481
+ if preds.ndim != 3 or preds.shape[0] != 1:
482
+ raise ValueError(f"Unexpected raw ONNX output shape: {preds.shape}")
483
+ preds = preds[0]
484
+ if preds.shape[0] <= 16 and preds.shape[1] > preds.shape[0]:
485
+ preds = preds.T
486
+ if preds.ndim != 2 or preds.shape[1] < 5:
487
+ raise ValueError(f"Unexpected raw output shape: {preds.shape}")
488
+
489
+ boxes_xywh = preds[:, :4].astype(np.float32)
490
+ cls_part = preds[:, 4:].astype(np.float32)
491
+ if cls_part.shape[1] == 1:
492
+ scores = cls_part[:, 0]
493
+ cls_ids = np.zeros(len(scores), dtype=np.int32)
494
+ else:
495
+ cls_ids = np.argmax(cls_part, axis=1).astype(np.int32)
496
+ scores = cls_part[np.arange(len(cls_part)), cls_ids]
497
+ cls_ids = self.cls_remap[cls_ids]
498
+
499
+ keep = self._conf_filter_mask(scores, cls_ids)
500
+ boxes_xywh = boxes_xywh[keep]
501
+ scores = scores[keep]
502
+ cls_ids = cls_ids[keep]
503
+ if len(boxes_xywh) == 0:
504
+ return []
505
+ boxes = self._xywh_to_xyxy(boxes_xywh)
506
+
507
+ pad_w, pad_h = pad
508
+ boxes[:, [0, 2]] -= pad_w
509
+ boxes[:, [1, 3]] -= pad_h
510
+ boxes /= ratio
511
+ boxes = self._clip_boxes(boxes, orig_size)
512
+
513
+ boxes, scores, cls_ids = self._per_view_pipeline(
514
+ boxes, scores, cls_ids, orig_size
515
+ )
516
+ return self._build_results(boxes, scores, cls_ids)
517
+
518
+ @staticmethod
519
+ def _build_results(boxes: np.ndarray, scores: np.ndarray,
520
+ cls_ids: np.ndarray) -> list[BoundingBox]:
521
+ results: list[BoundingBox] = []
522
+ for box, conf, cls_id in zip(boxes, scores, cls_ids):
523
+ x1, y1, x2, y2 = box.tolist()
524
+ if x2 <= x1 or y2 <= y1:
525
+ continue
526
+ results.append(
527
+ BoundingBox(
528
+ x1=int(math.floor(x1)),
529
+ y1=int(math.floor(y1)),
530
+ x2=int(math.ceil(x2)),
531
+ y2=int(math.ceil(y2)),
532
+ cls_id=int(cls_id),
533
+ conf=float(conf),
534
+ )
535
+ )
536
+ return results
537
+
538
+ def _postprocess(self, output: np.ndarray, ratio: float,
539
+ pad: tuple[float, float],
540
+ orig_size: tuple[int, int]) -> list[BoundingBox]:
541
+ if output.ndim == 2 and output.shape[1] >= 6:
542
+ return self._decode_final_dets(output, ratio, pad, orig_size)
543
+ if output.ndim == 3 and output.shape[0] == 1 and output.shape[2] == 6:
544
+ return self._decode_final_dets(output, ratio, pad, orig_size)
545
+ return self._decode_raw_yolo(output, ratio, pad, orig_size)
546
+
547
+ def _predict_single(self, image: np.ndarray) -> list[BoundingBox]:
548
+ if image is None:
549
+ raise ValueError("Input image is None")
550
+ if not isinstance(image, np.ndarray):
551
+ raise TypeError(f"Input is not numpy array: {type(image)}")
552
+ if image.ndim != 3:
553
+ raise ValueError(f"Expected HWC image, got shape={image.shape}")
554
+ if image.shape[2] != 3:
555
+ raise ValueError(f"Expected 3 channels, got shape={image.shape}")
556
+ if image.dtype != np.uint8:
557
+ image = image.astype(np.uint8)
558
+
559
+ input_tensor, ratio, pad, orig_size = self._preprocess(image)
560
+ expected = (1, 3, self.input_height, self.input_width)
561
+ if input_tensor.shape != expected:
562
+ raise ValueError(
563
+ f"Bad input tensor shape={input_tensor.shape}, expected={expected}"
564
+ )
565
+
566
+ outputs = self.session.run(self.output_names, {self.input_name: input_tensor})
567
+ return self._postprocess(outputs[0], ratio, pad, orig_size)
568
+
569
+ def _predict_tta(self, image: np.ndarray) -> list[BoundingBox]:
570
+ boxes_orig = self._predict_single(image)
571
+ flipped = cv2.flip(image, 1)
572
+ boxes_flip = self._predict_single(flipped)
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
+ all_boxes = boxes_orig + boxes_flip
582
+ if not all_boxes:
583
+ return []
584
+
585
+ coords = np.array(
586
+ [[b.x1, b.y1, b.x2, b.y2] for b in all_boxes], dtype=np.float32
587
+ )
588
+ scores = np.array([b.conf for b in all_boxes], dtype=np.float32)
589
+ cls_ids = np.array([b.cls_id for b in all_boxes], dtype=np.int32)
590
+
591
+ # Per-class hard NMS on the merged final-from-both-views set, then
592
+ # cluster-confidence boost. This step is independent of nms_mode.
593
+ hard_keep = self._per_class_hard_nms(coords, scores, cls_ids, self.iou_thres)
594
+ if len(hard_keep) == 0:
595
+ return []
596
+ if len(hard_keep) > self.max_det:
597
+ top = np.argsort(-scores[hard_keep])[: self.max_det]
598
+ hard_keep = hard_keep[top]
599
+ boosted = self._max_score_per_cluster(
600
+ coords[hard_keep], cls_ids[hard_keep],
601
+ coords, scores, cls_ids, self.iou_thres,
602
+ )
603
+
604
+ kept_coords = coords[hard_keep]
605
+ kept_cls = cls_ids[hard_keep]
606
+ if self.cross_class_dedup and len(kept_coords) > 1:
607
+ kept_coords, boosted, kept_cls = self._cross_class_dedup_op(
608
+ kept_coords, boosted, kept_cls, self.cross_iou_thresh
609
+ )
610
+
611
+ return [
612
+ BoundingBox(
613
+ x1=int(math.floor(kept_coords[j, 0])),
614
+ y1=int(math.floor(kept_coords[j, 1])),
615
+ x2=int(math.ceil(kept_coords[j, 2])),
616
+ y2=int(math.ceil(kept_coords[j, 3])),
617
+ cls_id=int(kept_cls[j]),
618
+ conf=float(boosted[j]),
619
+ )
620
+ for j in range(len(kept_coords))
621
+ ]
622
+
623
+ def predict_batch(self, batch_images: list[ndarray], offset: int,
624
+ n_keypoints: int) -> list[TVFrameResult]:
625
+ results: list[TVFrameResult] = []
626
+ for frame_number_in_batch, image in enumerate(batch_images):
627
+ try:
628
+ if self.use_tta:
629
+ boxes = self._predict_tta(image)
630
+ else:
631
+ boxes = self._predict_single(image)
632
+ except Exception as e:
633
+ print(f"Inference failed for frame {offset + frame_number_in_batch}: {e}")
634
+ boxes = []
635
+ results.append(
636
+ TVFrameResult(
637
+ frame_id=offset + frame_number_in_batch,
638
+ boxes=boxes,
639
+ keypoints=[(0, 0) for _ in range(max(0, int(n_keypoints)))],
640
+ )
641
+ )
642
+ return results
readme.md ADDED
@@ -0,0 +1 @@
 
 
1
+ ## YOLO26 ONNX detector
weights.onnx ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:34acdef12ab4e7a243bdc9badd6dc5d10da24bc2ca60836e7230100ec7ac115e
3
+ size 19298313