alfred8995 commited on
Commit
198e60a
·
verified ·
1 Parent(s): 6e7dad6

Upload folder using huggingface_hub

Browse files
Files changed (3) hide show
  1. chute_config.yml +23 -0
  2. miner.py +644 -0
  3. weights.onnx +3 -0
chute_config.yml ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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: 24
11
+ include:
12
+ - "a6000"
13
+ - "4090"
14
+ - "3090"
15
+ - "l4"
16
+ - "l40s"
17
+
18
+ Chute:
19
+ timeout_seconds: 900
20
+ concurrency: 4
21
+ max_instances: 5
22
+ scaling_threshold: 0.5
23
+ shutdown_after_seconds: 288000
miner.py ADDED
@@ -0,0 +1,644 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ def __init__(self,
28
+ path_hf_repo: Path,
29
+ conf_thres: float = 0.23,
30
+ conf_high: float = 0.63,
31
+ iou_thres: float = 0.68,
32
+ tta_match_iou: float = 0.71,
33
+ max_det: int = 85,
34
+ min_box_area: int = 325,
35
+ min_w: int = 5,
36
+ min_h: int = 3,
37
+ max_aspect_ratio: float = 4,
38
+ max_box_area_ratio: float = 0.72,
39
+ ) -> None:
40
+ model_path = path_hf_repo / "weights.onnx"
41
+
42
+ self.class_names = ['bus', 'car', 'truck', 'motorcycle']
43
+ model_class_order = ['car', 'bus', 'truck', 'motorcycle']
44
+ self._train_cls_to_canonical = np.array(
45
+ [self.class_names.index(n) for n in model_class_order],
46
+ dtype=np.int32
47
+ )
48
+ print("ORT version:", ort.__version__)
49
+
50
+ try:
51
+ ort.preload_dlls()
52
+ print("✅ onnxruntime.preload_dlls() success")
53
+ except Exception as e:
54
+ print(f"⚠️ preload_dlls failed: {e}")
55
+
56
+ print("ORT available providers BEFORE session:", ort.get_available_providers())
57
+
58
+ sess_options = ort.SessionOptions()
59
+ sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
60
+
61
+ try:
62
+ self.session = ort.InferenceSession(
63
+ str(model_path),
64
+ sess_options=sess_options,
65
+ providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
66
+ )
67
+ print("✅ Created ORT session with preferred CUDA provider list")
68
+ except Exception as e:
69
+ print(f"⚠️ CUDA session creation failed, falling back to CPU: {e}")
70
+ self.session = ort.InferenceSession(
71
+ str(model_path),
72
+ sess_options=sess_options,
73
+ providers=["CPUExecutionProvider"],
74
+ )
75
+
76
+ print("ORT session providers:", self.session.get_providers())
77
+
78
+ for inp in self.session.get_inputs():
79
+ print("INPUT:", inp.name, inp.shape, inp.type)
80
+
81
+ for out in self.session.get_outputs():
82
+ print("OUTPUT:", out.name, out.shape, out.type)
83
+
84
+ self.input_name = self.session.get_inputs()[0].name
85
+ self.output_names = [output.name for output in self.session.get_outputs()]
86
+ self.input_shape = self.session.get_inputs()[0].shape
87
+
88
+ self.input_height = self._safe_dim(self.input_shape[2], default=1280)
89
+ self.input_width = self._safe_dim(self.input_shape[3], default=1280)
90
+
91
+ # ---------- Scoring-oriented thresholds ----------
92
+ # Low threshold for candidate generation
93
+ self.conf_thres = conf_thres
94
+
95
+ # High-confidence boxes can survive without TTA confirmation
96
+ self.conf_high = conf_high
97
+
98
+ # NMS threshold
99
+ self.iou_thres = iou_thres
100
+
101
+ # TTA confirmation IoU
102
+ self.tta_match_iou = tta_match_iou
103
+
104
+ self.max_det = max_det
105
+ self.use_tta = True
106
+
107
+ # Box sanity filters
108
+ self.min_box_area = min_box_area
109
+ self.min_w = min_w
110
+ self.min_h = min_h
111
+ self.max_aspect_ratio = max_aspect_ratio
112
+ self.max_box_area_ratio = max_box_area_ratio
113
+
114
+ print(f"✅ ONNX model loaded from: {model_path}")
115
+ print(f"✅ ONNX providers: {self.session.get_providers()}")
116
+ print(f"✅ ONNX input: name={self.input_name}, shape={self.input_shape}")
117
+
118
+ def __repr__(self) -> str:
119
+ return (
120
+ f"ONNXRuntime(session={type(self.session).__name__}, "
121
+ f"providers={self.session.get_providers()})"
122
+ )
123
+
124
+ @staticmethod
125
+ def _safe_dim(value, default: int) -> int:
126
+ return value if isinstance(value, int) and value > 0 else default
127
+
128
+ def _remap_train_cls_ids(self, cls_ids: np.ndarray) -> np.ndarray:
129
+ idx = np.clip(cls_ids.astype(np.int64, copy=False), 0, len(self._train_cls_to_canonical) - 1)
130
+ return self._train_cls_to_canonical[idx]
131
+
132
+ def _letterbox(
133
+ self,
134
+ image: ndarray,
135
+ new_shape: tuple[int, int],
136
+ color=(114, 114, 114),
137
+ ) -> tuple[ndarray, float, tuple[float, float]]:
138
+ h, w = image.shape[:2]
139
+ new_w, new_h = new_shape
140
+
141
+ ratio = min(new_w / w, new_h / h)
142
+ resized_w = int(round(w * ratio))
143
+ resized_h = int(round(h * ratio))
144
+
145
+ if (resized_w, resized_h) != (w, h):
146
+ interp = cv2.INTER_CUBIC if ratio > 1.0 else cv2.INTER_LINEAR
147
+ image = cv2.resize(image, (resized_w, resized_h), interpolation=interp)
148
+
149
+ dw = new_w - resized_w
150
+ dh = new_h - resized_h
151
+ dw /= 2.0
152
+ dh /= 2.0
153
+
154
+ left = int(round(dw - 0.1))
155
+ right = int(round(dw + 0.1))
156
+ top = int(round(dh - 0.1))
157
+ bottom = int(round(dh + 0.1))
158
+
159
+ padded = cv2.copyMakeBorder(
160
+ image,
161
+ top,
162
+ bottom,
163
+ left,
164
+ right,
165
+ borderType=cv2.BORDER_CONSTANT,
166
+ value=color,
167
+ )
168
+ return padded, ratio, (dw, dh)
169
+
170
+ def _preprocess(
171
+ self, image: ndarray
172
+ ) -> tuple[np.ndarray, float, tuple[float, float], tuple[int, int]]:
173
+ orig_h, orig_w = image.shape[:2]
174
+
175
+ img, ratio, pad = self._letterbox(
176
+ image, (self.input_width, self.input_height)
177
+ )
178
+ img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
179
+ img = img.astype(np.float32) / 255.0
180
+ img = np.transpose(img, (2, 0, 1))[None, ...]
181
+ img = np.ascontiguousarray(img, dtype=np.float32)
182
+
183
+ return img, ratio, pad, (orig_w, orig_h)
184
+
185
+ @staticmethod
186
+ def _clip_boxes(boxes: np.ndarray, image_size: tuple[int, int]) -> np.ndarray:
187
+ w, h = image_size
188
+ boxes[:, 0] = np.clip(boxes[:, 0], 0, w - 1)
189
+ boxes[:, 1] = np.clip(boxes[:, 1], 0, h - 1)
190
+ boxes[:, 2] = np.clip(boxes[:, 2], 0, w - 1)
191
+ boxes[:, 3] = np.clip(boxes[:, 3], 0, h - 1)
192
+ return boxes
193
+
194
+ @staticmethod
195
+ def _xywh_to_xyxy(boxes: np.ndarray) -> np.ndarray:
196
+ out = np.empty_like(boxes)
197
+ out[:, 0] = boxes[:, 0] - boxes[:, 2] / 2.0
198
+ out[:, 1] = boxes[:, 1] - boxes[:, 3] / 2.0
199
+ out[:, 2] = boxes[:, 0] + boxes[:, 2] / 2.0
200
+ out[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.0
201
+ return out
202
+
203
+ @staticmethod
204
+ def _hard_nms(
205
+ boxes: np.ndarray,
206
+ scores: np.ndarray,
207
+ iou_thresh: float,
208
+ ) -> np.ndarray:
209
+ if len(boxes) == 0:
210
+ return np.array([], dtype=np.intp)
211
+
212
+ boxes = np.asarray(boxes, dtype=np.float32)
213
+ scores = np.asarray(scores, dtype=np.float32)
214
+ order = np.argsort(scores)[::-1]
215
+ keep = []
216
+
217
+ while len(order) > 0:
218
+ i = order[0]
219
+ keep.append(i)
220
+ if len(order) == 1:
221
+ break
222
+
223
+ rest = order[1:]
224
+
225
+ xx1 = np.maximum(boxes[i, 0], boxes[rest, 0])
226
+ yy1 = np.maximum(boxes[i, 1], boxes[rest, 1])
227
+ xx2 = np.minimum(boxes[i, 2], boxes[rest, 2])
228
+ yy2 = np.minimum(boxes[i, 3], boxes[rest, 3])
229
+
230
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
231
+
232
+ area_i = np.maximum(0.0, (boxes[i, 2] - boxes[i, 0])) * np.maximum(0.0, (boxes[i, 3] - boxes[i, 1]))
233
+ area_r = np.maximum(0.0, (boxes[rest, 2] - boxes[rest, 0])) * np.maximum(0.0, (boxes[rest, 3] - boxes[rest, 1]))
234
+
235
+ iou = inter / (area_i + area_r - inter + 1e-7)
236
+ order = rest[iou <= iou_thresh]
237
+
238
+ return np.array(keep, dtype=np.intp)
239
+
240
+ @classmethod
241
+ def _nms_per_class(
242
+ cls,
243
+ boxes: np.ndarray,
244
+ scores: np.ndarray,
245
+ cls_ids: np.ndarray,
246
+ iou_thresh: float,
247
+ max_det: int,
248
+ ) -> np.ndarray:
249
+ """NMS within each class so overlapping car vs bus predictions are not merged away."""
250
+ if len(boxes) == 0:
251
+ return np.array([], dtype=np.intp)
252
+ keep_all: list[int] = []
253
+ for c in np.unique(cls_ids):
254
+ idxs = np.nonzero(cls_ids == c)[0]
255
+ if len(idxs) == 0:
256
+ continue
257
+ local_keep = cls._hard_nms(boxes[idxs], scores[idxs], iou_thresh)
258
+ keep_all.extend(idxs[local_keep].tolist())
259
+ keep_all = np.array(keep_all, dtype=np.intp)
260
+ order = np.argsort(scores[keep_all])[::-1]
261
+ return keep_all[order[:max_det]]
262
+
263
+ @staticmethod
264
+ def _box_iou_one_to_many(box: np.ndarray, boxes: np.ndarray) -> np.ndarray:
265
+ xx1 = np.maximum(box[0], boxes[:, 0])
266
+ yy1 = np.maximum(box[1], boxes[:, 1])
267
+ xx2 = np.minimum(box[2], boxes[:, 2])
268
+ yy2 = np.minimum(box[3], boxes[:, 3])
269
+
270
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
271
+
272
+ area_a = max(0.0, (box[2] - box[0]) * (box[3] - box[1]))
273
+ area_b = np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) * np.maximum(0.0, boxes[:, 3] - boxes[:, 1])
274
+
275
+ return inter / (area_a + area_b - inter + 1e-7)
276
+
277
+ def _filter_sane_boxes(
278
+ self,
279
+ boxes: np.ndarray,
280
+ scores: np.ndarray,
281
+ cls_ids: np.ndarray,
282
+ orig_size: tuple[int, int],
283
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
284
+ if len(boxes) == 0:
285
+ return boxes, scores, cls_ids
286
+
287
+ orig_w, orig_h = orig_size
288
+ image_area = float(orig_w * orig_h)
289
+
290
+ keep = []
291
+ for i, box in enumerate(boxes):
292
+ x1, y1, x2, y2 = box.tolist()
293
+ bw = x2 - x1
294
+ bh = y2 - y1
295
+
296
+ if bw <= 0 or bh <= 0:
297
+ continue
298
+ if bw < self.min_w or bh < self.min_h:
299
+ continue
300
+
301
+ area = bw * bh
302
+ if area < self.min_box_area:
303
+ continue
304
+ if area > self.max_box_area_ratio * image_area:
305
+ continue
306
+
307
+ ar = max(bw / max(bh, 1e-6), bh / max(bw, 1e-6))
308
+ if ar > self.max_aspect_ratio:
309
+ continue
310
+
311
+ keep.append(i)
312
+
313
+ if not keep:
314
+ return (
315
+ np.empty((0, 4), dtype=np.float32),
316
+ np.empty((0,), dtype=np.float32),
317
+ np.empty((0,), dtype=np.int32),
318
+ )
319
+
320
+ keep = np.array(keep, dtype=np.intp)
321
+ return boxes[keep], scores[keep], cls_ids[keep]
322
+
323
+ def _decode_final_dets(
324
+ self,
325
+ preds: np.ndarray,
326
+ ratio: float,
327
+ pad: tuple[float, float],
328
+ orig_size: tuple[int, int],
329
+ ) -> list[BoundingBox]:
330
+ if preds.ndim == 3 and preds.shape[0] == 1:
331
+ preds = preds[0]
332
+
333
+ if preds.ndim != 2 or preds.shape[1] < 6:
334
+ raise ValueError(f"Unexpected ONNX final-det output shape: {preds.shape}")
335
+
336
+ boxes = preds[:, :4].astype(np.float32)
337
+ scores = preds[:, 4].astype(np.float32)
338
+ cls_ids = self._remap_train_cls_ids(preds[:, 5].astype(np.int32))
339
+ # All trained vehicle classes: bus, car, truck, motorcycle (see self.class_names).
340
+
341
+ # candidate threshold
342
+ keep = scores >= self.conf_thres
343
+ boxes = boxes[keep]
344
+ scores = scores[keep]
345
+ cls_ids = cls_ids[keep]
346
+
347
+ if len(boxes) == 0:
348
+ return []
349
+
350
+ pad_w, pad_h = pad
351
+ orig_w, orig_h = orig_size
352
+
353
+ boxes[:, [0, 2]] -= pad_w
354
+ boxes[:, [1, 3]] -= pad_h
355
+ boxes /= ratio
356
+ boxes = self._clip_boxes(boxes, (orig_w, orig_h))
357
+
358
+ boxes, scores, cls_ids = self._filter_sane_boxes(boxes, scores, cls_ids, orig_size)
359
+ if len(boxes) == 0:
360
+ return []
361
+
362
+ keep_idx = self._nms_per_class(
363
+ boxes, scores, cls_ids, self.iou_thres, self.max_det
364
+ )
365
+
366
+ boxes = boxes[keep_idx]
367
+ scores = scores[keep_idx]
368
+ cls_ids = cls_ids[keep_idx]
369
+
370
+ return [
371
+ BoundingBox(
372
+ x1=int(math.floor(box[0])),
373
+ y1=int(math.floor(box[1])),
374
+ x2=int(math.ceil(box[2])),
375
+ y2=int(math.ceil(box[3])),
376
+ cls_id=int(cls_id),
377
+ conf=float(conf),
378
+ )
379
+ for box, conf, cls_id in zip(boxes, scores, cls_ids)
380
+ if box[2] > box[0] and box[3] > box[1]
381
+ ]
382
+
383
+ def _decode_raw_yolo(
384
+ self,
385
+ preds: np.ndarray,
386
+ ratio: float,
387
+ pad: tuple[float, float],
388
+ orig_size: tuple[int, int],
389
+ ) -> list[BoundingBox]:
390
+ if preds.ndim != 3:
391
+ raise ValueError(f"Unexpected raw ONNX output shape: {preds.shape}")
392
+ if preds.shape[0] != 1:
393
+ raise ValueError(f"Unexpected batch dimension in raw output: {preds.shape}")
394
+
395
+ preds = preds[0]
396
+
397
+ # Normalize to [N, C]
398
+ if preds.shape[0] <= 16 and preds.shape[1] > preds.shape[0]:
399
+ preds = preds.T
400
+
401
+ if preds.ndim != 2 or preds.shape[1] < 5:
402
+ raise ValueError(f"Unexpected normalized raw output shape: {preds.shape}")
403
+
404
+ boxes_xywh = preds[:, :4].astype(np.float32)
405
+ tail = preds[:, 4:].astype(np.float32)
406
+
407
+ # Supports:
408
+ # [x,y,w,h,score] single-class
409
+ # [x,y,w,h,obj,cls] YOLO standard single-class
410
+ # [x,y,w,h,obj,cls1,cls2,...] multi-class
411
+ if tail.shape[1] == 1:
412
+ scores = tail[:, 0]
413
+ cls_ids = np.zeros(len(scores), dtype=np.int32)
414
+ elif tail.shape[1] == 2:
415
+ obj = tail[:, 0]
416
+ cls_prob = tail[:, 1]
417
+ scores = obj * cls_prob
418
+ cls_ids = np.zeros(len(scores), dtype=np.int32)
419
+ else:
420
+ obj = tail[:, 0]
421
+ class_probs = tail[:, 1:]
422
+ cls_ids = np.argmax(class_probs, axis=1).astype(np.int32)
423
+ cls_scores = class_probs[np.arange(len(class_probs)), cls_ids]
424
+ scores = obj * cls_scores
425
+
426
+ cls_ids = self._remap_train_cls_ids(cls_ids)
427
+
428
+ keep = scores >= self.conf_thres
429
+ boxes_xywh = boxes_xywh[keep]
430
+ scores = scores[keep]
431
+ cls_ids = cls_ids[keep]
432
+
433
+ if len(boxes_xywh) == 0:
434
+ return []
435
+
436
+ boxes = self._xywh_to_xyxy(boxes_xywh)
437
+
438
+ pad_w, pad_h = pad
439
+ orig_w, orig_h = orig_size
440
+
441
+ boxes[:, [0, 2]] -= pad_w
442
+ boxes[:, [1, 3]] -= pad_h
443
+ boxes /= ratio
444
+ boxes = self._clip_boxes(boxes, (orig_w, orig_h))
445
+
446
+ boxes, scores, cls_ids = self._filter_sane_boxes(boxes, scores, cls_ids, orig_size)
447
+ if len(boxes) == 0:
448
+ return []
449
+
450
+ keep_idx = self._nms_per_class(
451
+ boxes, scores, cls_ids, self.iou_thres, self.max_det
452
+ )
453
+
454
+ boxes = boxes[keep_idx]
455
+ scores = scores[keep_idx]
456
+ cls_ids = cls_ids[keep_idx]
457
+
458
+ return [
459
+ BoundingBox(
460
+ x1=int(math.floor(box[0])),
461
+ y1=int(math.floor(box[1])),
462
+ x2=int(math.ceil(box[2])),
463
+ y2=int(math.ceil(box[3])),
464
+ cls_id=int(cls_id),
465
+ conf=float(conf),
466
+ )
467
+ for box, conf, cls_id in zip(boxes, scores, cls_ids)
468
+ if box[2] > box[0] and box[3] > box[1]
469
+ ]
470
+
471
+ def _postprocess(
472
+ self,
473
+ output: np.ndarray,
474
+ ratio: float,
475
+ pad: tuple[float, float],
476
+ orig_size: tuple[int, int],
477
+ ) -> list[BoundingBox]:
478
+ if output.ndim == 2 and output.shape[1] >= 6:
479
+ return self._decode_final_dets(output, ratio, pad, orig_size)
480
+
481
+ if output.ndim == 3 and output.shape[0] == 1 and output.shape[2] >= 6:
482
+ return self._decode_final_dets(output, ratio, pad, orig_size)
483
+
484
+ return self._decode_raw_yolo(output, ratio, pad, orig_size)
485
+
486
+ def _predict_single(self, image: np.ndarray) -> list[BoundingBox]:
487
+ if image is None:
488
+ raise ValueError("Input image is None")
489
+ if not isinstance(image, np.ndarray):
490
+ raise TypeError(f"Input is not numpy array: {type(image)}")
491
+ if image.ndim != 3:
492
+ raise ValueError(f"Expected HWC image, got shape={image.shape}")
493
+ if image.shape[0] <= 0 or image.shape[1] <= 0:
494
+ raise ValueError(f"Invalid image shape={image.shape}")
495
+ if image.shape[2] != 3:
496
+ raise ValueError(f"Expected 3 channels, got shape={image.shape}")
497
+
498
+ if image.dtype != np.uint8:
499
+ image = image.astype(np.uint8)
500
+
501
+ input_tensor, ratio, pad, orig_size = self._preprocess(image)
502
+
503
+ expected_shape = (1, 3, self.input_height, self.input_width)
504
+ if input_tensor.shape != expected_shape:
505
+ raise ValueError(
506
+ f"Bad input tensor shape={input_tensor.shape}, expected={expected_shape}"
507
+ )
508
+
509
+ outputs = self.session.run(self.output_names, {self.input_name: input_tensor})
510
+ det_output = outputs[0]
511
+ return self._postprocess(det_output, ratio, pad, orig_size)
512
+
513
+ def _merge_tta_consensus(
514
+ self,
515
+ boxes_orig: list[BoundingBox],
516
+ boxes_flip: list[BoundingBox],
517
+ ) -> list[BoundingBox]:
518
+ """
519
+ Keep:
520
+ - any box with conf >= conf_high
521
+ - low/medium-conf boxes only if confirmed across TTA views
522
+ Then run final hard NMS.
523
+ """
524
+ if not boxes_orig and not boxes_flip:
525
+ return []
526
+
527
+ coords_o = np.array([[b.x1, b.y1, b.x2, b.y2] for b in boxes_orig], dtype=np.float32) if boxes_orig else np.empty((0, 4), dtype=np.float32)
528
+ scores_o = np.array([b.conf for b in boxes_orig], dtype=np.float32) if boxes_orig else np.empty((0,), dtype=np.float32)
529
+ cls_o = np.array([b.cls_id for b in boxes_orig], dtype=np.int32) if boxes_orig else np.empty((0,), dtype=np.int32)
530
+
531
+ coords_f = np.array([[b.x1, b.y1, b.x2, b.y2] for b in boxes_flip], dtype=np.float32) if boxes_flip else np.empty((0, 4), dtype=np.float32)
532
+ scores_f = np.array([b.conf for b in boxes_flip], dtype=np.float32) if boxes_flip else np.empty((0,), dtype=np.float32)
533
+ cls_f = np.array([b.cls_id for b in boxes_flip], dtype=np.int32) if boxes_flip else np.empty((0,), dtype=np.int32)
534
+
535
+ accepted_boxes = []
536
+ accepted_scores = []
537
+ accepted_cls = []
538
+
539
+ # Original view candidates
540
+ for i in range(len(coords_o)):
541
+ score = scores_o[i]
542
+ if score >= self.conf_high:
543
+ accepted_boxes.append(coords_o[i])
544
+ accepted_scores.append(score)
545
+ accepted_cls.append(int(cls_o[i]))
546
+ elif len(coords_f) > 0:
547
+ ious = self._box_iou_one_to_many(coords_o[i], coords_f)
548
+ j = int(np.argmax(ious))
549
+ if ious[j] >= self.tta_match_iou:
550
+ fused_score = max(score, scores_f[j])
551
+ accepted_boxes.append(coords_o[i])
552
+ accepted_scores.append(fused_score)
553
+ accepted_cls.append(int(cls_o[i]))
554
+
555
+ # Flipped-view high-confidence boxes that original missed
556
+ for i in range(len(coords_f)):
557
+ score = scores_f[i]
558
+ if score < self.conf_high:
559
+ continue
560
+
561
+ if len(coords_o) == 0:
562
+ accepted_boxes.append(coords_f[i])
563
+ accepted_scores.append(score)
564
+ accepted_cls.append(int(cls_f[i]))
565
+ continue
566
+
567
+ ious = self._box_iou_one_to_many(coords_f[i], coords_o)
568
+ if np.max(ious) < self.tta_match_iou:
569
+ accepted_boxes.append(coords_f[i])
570
+ accepted_scores.append(score)
571
+ accepted_cls.append(int(cls_f[i]))
572
+
573
+ if not accepted_boxes:
574
+ return []
575
+
576
+ boxes = np.array(accepted_boxes, dtype=np.float32)
577
+ scores = np.array(accepted_scores, dtype=np.float32)
578
+ cls_ids = np.array(accepted_cls, dtype=np.int32)
579
+
580
+ keep = self._nms_per_class(boxes, scores, cls_ids, self.iou_thres, self.max_det)
581
+
582
+ out = []
583
+ for idx in keep:
584
+ x1, y1, x2, y2 = boxes[idx].tolist()
585
+ out.append(
586
+ BoundingBox(
587
+ x1=int(math.floor(x1)),
588
+ y1=int(math.floor(y1)),
589
+ x2=int(math.ceil(x2)),
590
+ y2=int(math.ceil(y2)),
591
+ cls_id=int(cls_ids[idx]),
592
+ conf=float(scores[idx]),
593
+ )
594
+ )
595
+ return out
596
+
597
+ def _predict_tta(self, image: np.ndarray) -> list[BoundingBox]:
598
+ boxes_orig = self._predict_single(image)
599
+
600
+ flipped = cv2.flip(image, 1)
601
+ boxes_flip_raw = self._predict_single(flipped)
602
+
603
+ w = image.shape[1]
604
+ boxes_flip = [
605
+ BoundingBox(
606
+ x1=w - b.x2,
607
+ y1=b.y1,
608
+ x2=w - b.x1,
609
+ y2=b.y2,
610
+ cls_id=b.cls_id,
611
+ conf=b.conf,
612
+ )
613
+ for b in boxes_flip_raw
614
+ ]
615
+
616
+ return self._merge_tta_consensus(boxes_orig, boxes_flip)
617
+
618
+ def predict_batch(
619
+ self,
620
+ batch_images: list[ndarray],
621
+ offset: int,
622
+ n_keypoints: int,
623
+ ) -> list[TVFrameResult]:
624
+ results: list[TVFrameResult] = []
625
+
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
+
636
+ results.append(
637
+ TVFrameResult(
638
+ frame_id=offset + frame_number_in_batch,
639
+ boxes=boxes,
640
+ keypoints=[(0, 0) for _ in range(max(0, int(n_keypoints)))],
641
+ )
642
+ )
643
+
644
+ return results
weights.onnx ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:b3abbd023f1daf0fbf858abf6f47c6f2d294d5e37486b2f62df1e21b164827e1
3
+ size 19407982