Realfencer commited on
Commit
d7c189e
·
verified ·
1 Parent(s): 99c377a

Upload 3 files

Browse files
Files changed (3) hide show
  1. chute_config.yml +21 -0
  2. miner.py +528 -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,528 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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. Hard global NMS + sanity filter + dedup + flip TTA, with per-class rescue bonus."""
28
+
29
+ class_names = ["cup", "bottle", "can"]
30
+ model_class_names = ["bottle", "can", "cup"]
31
+ _model_to_competition_cls = np.array([1, 2, 0], dtype=np.int32)
32
+ input_size = 1280
33
+ iou_thres = 0.4
34
+ cross_iou_thresh = 0.7
35
+ min_side = 8.0
36
+ min_box_area = 100.0
37
+ max_aspect_ratio = 10.0
38
+ max_det = 300
39
+ _conf_thres_array = np.array([0.6, 0.45, 0.5], dtype=np.float32)
40
+ _bonus_array = np.array([0.0, 0.0, 0.2], dtype=np.float32)
41
+
42
+ def __init__(self, path_hf_repo: Path) -> None:
43
+ model_path = path_hf_repo / "weights.onnx"
44
+ print("ORT version:", ort.__version__)
45
+
46
+ try:
47
+ ort.preload_dlls()
48
+ print("preload_dlls success")
49
+ except Exception as e:
50
+ print(f"preload_dlls failed: {e}")
51
+
52
+ print("ORT available providers BEFORE session:", ort.get_available_providers())
53
+
54
+ sess_options = ort.SessionOptions()
55
+ sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
56
+
57
+ try:
58
+ self.session = ort.InferenceSession(
59
+ str(model_path),
60
+ sess_options=sess_options,
61
+ providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
62
+ )
63
+ print("Created ORT session with preferred CUDA provider list")
64
+ except Exception as e:
65
+ print(f"CUDA session creation failed, falling back to CPU: {e}")
66
+ self.session = ort.InferenceSession(
67
+ str(model_path),
68
+ sess_options=sess_options,
69
+ providers=["CPUExecutionProvider"],
70
+ )
71
+
72
+ print("ORT session providers:", self.session.get_providers())
73
+
74
+ for inp in self.session.get_inputs():
75
+ print("INPUT:", inp.name, inp.shape, inp.type)
76
+ for out in self.session.get_outputs():
77
+ print("OUTPUT:", out.name, out.shape, out.type)
78
+
79
+ self.input_name = self.session.get_inputs()[0].name
80
+ self.output_names = [output.name for output in self.session.get_outputs()]
81
+ self.input_shape = self.session.get_inputs()[0].shape
82
+ self.input_dtype = self._input_dtype(self.session.get_inputs()[0].type)
83
+
84
+ self.input_height = self._safe_dim(self.input_shape[2], default=self.input_size)
85
+ self.input_width = self._safe_dim(self.input_shape[3], default=self.input_size)
86
+
87
+ print(f"ONNX model loaded from: {model_path}")
88
+ print(f"ONNX providers: {self.session.get_providers()}")
89
+ print(f"ONNX input: name={self.input_name}, shape={self.input_shape}")
90
+ print(f"ONNX input dtype: {self.input_dtype}")
91
+
92
+ def __repr__(self) -> str:
93
+ return (
94
+ f"ONNXRuntime(session={type(self.session).__name__}, "
95
+ f"providers={self.session.get_providers()})"
96
+ )
97
+
98
+ @staticmethod
99
+ def _safe_dim(value, default: int) -> int:
100
+ return value if isinstance(value, int) and value > 0 else default
101
+
102
+ @staticmethod
103
+ def _input_dtype(input_type: str) -> np.dtype:
104
+ if input_type == "tensor(float16)":
105
+ return np.dtype(np.float16)
106
+ return np.dtype(np.float32)
107
+
108
+ @classmethod
109
+ def _to_competition_cls(cls, cls_ids: np.ndarray) -> np.ndarray:
110
+ if len(cls_ids) == 0:
111
+ return cls_ids.astype(np.int32)
112
+ valid = (cls_ids >= 0) & (cls_ids < len(cls._model_to_competition_cls))
113
+ remapped = np.full_like(cls_ids, fill_value=-1, dtype=np.int32)
114
+ remapped[valid] = cls._model_to_competition_cls[cls_ids[valid]]
115
+ return remapped
116
+
117
+ def _letterbox(self, image: ndarray, new_shape: tuple[int, int],
118
+ color=(114, 114, 114)
119
+ ) -> tuple[ndarray, float, tuple[float, float]]:
120
+ h, w = image.shape[:2]
121
+ new_w, new_h = new_shape
122
+ ratio = min(new_w / w, new_h / h)
123
+ resized_w = int(round(w * ratio))
124
+ resized_h = int(round(h * ratio))
125
+ if (resized_w, resized_h) != (w, h):
126
+ interp = cv2.INTER_CUBIC if ratio > 1.0 else cv2.INTER_LINEAR
127
+ image = cv2.resize(image, (resized_w, resized_h), interpolation=interp)
128
+ dw = (new_w - resized_w) / 2.0
129
+ dh = (new_h - resized_h) / 2.0
130
+ left = int(round(dw - 0.1))
131
+ right = int(round(dw + 0.1))
132
+ top = int(round(dh - 0.1))
133
+ bottom = int(round(dh + 0.1))
134
+ padded = cv2.copyMakeBorder(image, top, bottom, left, right,
135
+ borderType=cv2.BORDER_CONSTANT, value=color)
136
+ return padded, ratio, (dw, dh)
137
+
138
+ def _preprocess(self, image: ndarray
139
+ ) -> tuple[np.ndarray, float, tuple[float, float],
140
+ tuple[int, int]]:
141
+ orig_h, orig_w = image.shape[:2]
142
+ img, ratio, pad = self._letterbox(image, (self.input_width, self.input_height))
143
+ img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
144
+ img = img.astype(self.input_dtype) / self.input_dtype.type(255.0)
145
+ img = np.transpose(img, (2, 0, 1))[None, ...]
146
+ img = np.ascontiguousarray(img, dtype=self.input_dtype)
147
+ return img, ratio, pad, (orig_w, orig_h)
148
+
149
+ @staticmethod
150
+ def _clip_boxes(boxes: np.ndarray, image_size: tuple[int, int]) -> np.ndarray:
151
+ w, h = image_size
152
+ boxes[:, 0] = np.clip(boxes[:, 0], 0, w - 1)
153
+ boxes[:, 1] = np.clip(boxes[:, 1], 0, h - 1)
154
+ boxes[:, 2] = np.clip(boxes[:, 2], 0, w - 1)
155
+ boxes[:, 3] = np.clip(boxes[:, 3], 0, h - 1)
156
+ return boxes
157
+
158
+ @staticmethod
159
+ def _xywh_to_xyxy(boxes: np.ndarray) -> np.ndarray:
160
+ out = np.empty_like(boxes)
161
+ out[:, 0] = boxes[:, 0] - boxes[:, 2] / 2.0
162
+ out[:, 1] = boxes[:, 1] - boxes[:, 3] / 2.0
163
+ out[:, 2] = boxes[:, 0] + boxes[:, 2] / 2.0
164
+ out[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.0
165
+ return out
166
+
167
+ @staticmethod
168
+ def _hard_nms(boxes: np.ndarray, scores: np.ndarray,
169
+ iou_thresh: float) -> np.ndarray:
170
+ n = len(boxes)
171
+ if n == 0:
172
+ return np.array([], dtype=np.intp)
173
+ order = np.argsort(-scores)
174
+ keep: list[int] = []
175
+ while len(order) > 0:
176
+ i = int(order[0])
177
+ keep.append(i)
178
+ if len(order) == 1:
179
+ break
180
+ rest = order[1:]
181
+ xx1 = np.maximum(boxes[i, 0], boxes[rest, 0])
182
+ yy1 = np.maximum(boxes[i, 1], boxes[rest, 1])
183
+ xx2 = np.minimum(boxes[i, 2], boxes[rest, 2])
184
+ yy2 = np.minimum(boxes[i, 3], boxes[rest, 3])
185
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
186
+ a_i = (max(0.0, boxes[i, 2] - boxes[i, 0]) *
187
+ max(0.0, boxes[i, 3] - boxes[i, 1]))
188
+ a_r = (np.maximum(0.0, boxes[rest, 2] - boxes[rest, 0]) *
189
+ np.maximum(0.0, boxes[rest, 3] - boxes[rest, 1]))
190
+ iou = inter / (a_i + a_r - inter + 1e-7)
191
+ order = rest[iou <= iou_thresh]
192
+ return np.array(keep, dtype=np.intp)
193
+
194
+ def _per_class_hard_nms(self, boxes: np.ndarray, scores: np.ndarray,
195
+ cls_ids: np.ndarray, iou_thresh: float
196
+ ) -> np.ndarray:
197
+ if len(boxes) == 0:
198
+ return np.array([], dtype=np.intp)
199
+ all_keep: list[int] = []
200
+ for c in np.unique(cls_ids):
201
+ mask = cls_ids == c
202
+ indices = np.where(mask)[0]
203
+ keep = self._hard_nms(boxes[mask], scores[mask], iou_thresh)
204
+ all_keep.extend(indices[keep].tolist())
205
+ all_keep.sort()
206
+ return np.array(all_keep, dtype=np.intp)
207
+
208
+ def _cross_class_dedup_op(self, boxes: np.ndarray, scores: np.ndarray,
209
+ cls_ids: np.ndarray, iou_thresh: float
210
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
211
+ n = len(boxes)
212
+ if n <= 1:
213
+ return boxes, scores, cls_ids
214
+ boxes = np.asarray(boxes, dtype=np.float32)
215
+ scores = np.asarray(scores, dtype=np.float32)
216
+ cls_ids = np.asarray(cls_ids, dtype=np.int32)
217
+ areas = (np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) *
218
+ np.maximum(0.0, boxes[:, 3] - boxes[:, 1]))
219
+ margins = scores - self._conf_thres_array[cls_ids]
220
+ order = np.lexsort((-areas, -margins))
221
+ suppressed = np.zeros(n, dtype=bool)
222
+ keep: list[int] = []
223
+ for i in order:
224
+ if suppressed[i]:
225
+ continue
226
+ keep.append(int(i))
227
+ bi = boxes[i]
228
+ xx1 = np.maximum(bi[0], boxes[:, 0])
229
+ yy1 = np.maximum(bi[1], boxes[:, 1])
230
+ xx2 = np.minimum(bi[2], boxes[:, 2])
231
+ yy2 = np.minimum(bi[3], boxes[:, 3])
232
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
233
+ a_i = max(1e-7, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
234
+ iou = inter / (a_i + areas - inter + 1e-7)
235
+ dup = iou > iou_thresh
236
+ dup[i] = False
237
+ suppressed |= dup
238
+ keep_idx = np.array(keep, dtype=np.intp)
239
+ return boxes[keep_idx], scores[keep_idx], cls_ids[keep_idx]
240
+
241
+ def _filter_sane_boxes(self, boxes: np.ndarray, scores: np.ndarray,
242
+ cls_ids: np.ndarray, orig_size: tuple[int, int]
243
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
244
+ if len(boxes) == 0:
245
+ return boxes, scores, cls_ids
246
+ orig_w, orig_h = orig_size
247
+ image_area = float(orig_w * orig_h)
248
+ bw = np.maximum(0.0, boxes[:, 2] - boxes[:, 0])
249
+ bh = np.maximum(0.0, boxes[:, 3] - boxes[:, 1])
250
+ area = bw * bh
251
+ ar = np.where(
252
+ (bw > 0) & (bh > 0),
253
+ np.maximum(bw / np.maximum(bh, 1e-6), bh / np.maximum(bw, 1e-6)),
254
+ np.inf,
255
+ )
256
+ keep = (
257
+ (bw >= self.min_side) & (bh >= self.min_side) &
258
+ (area >= self.min_box_area) &
259
+ (area <= 0.95 * image_area) &
260
+ (ar <= self.max_aspect_ratio)
261
+ )
262
+ return boxes[keep], scores[keep], cls_ids[keep]
263
+
264
+ def _max_score_per_cluster(self, post_boxes: np.ndarray,
265
+ post_cls: np.ndarray,
266
+ full_boxes: np.ndarray,
267
+ full_scores: np.ndarray,
268
+ full_cls: np.ndarray,
269
+ iou_thresh: float) -> np.ndarray:
270
+ n = len(post_boxes)
271
+ if n == 0:
272
+ return np.empty(0, dtype=np.float32)
273
+ full_areas = (np.maximum(0.0, full_boxes[:, 2] - full_boxes[:, 0]) *
274
+ np.maximum(0.0, full_boxes[:, 3] - full_boxes[:, 1]))
275
+ out = np.empty(n, dtype=np.float32)
276
+ for i in range(n):
277
+ bi = post_boxes[i]
278
+ xx1 = np.maximum(bi[0], full_boxes[:, 0])
279
+ yy1 = np.maximum(bi[1], full_boxes[:, 1])
280
+ xx2 = np.minimum(bi[2], full_boxes[:, 2])
281
+ yy2 = np.minimum(bi[3], full_boxes[:, 3])
282
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
283
+ a_i = max(0.0, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
284
+ iou = inter / (a_i + full_areas - inter + 1e-7)
285
+ cluster = (iou >= iou_thresh) & (full_cls == post_cls[i])
286
+ out[i] = float(np.max(full_scores[cluster])) if np.any(cluster) else 0.0
287
+ return out
288
+
289
+ def _conf_filter_mask(self, scores: np.ndarray,
290
+ cls_ids: np.ndarray) -> np.ndarray:
291
+ """Boolean keep-mask: score >= per-class threshold, with a per-class
292
+ rescue — if a class has zero boxes passing, admit its top-1 candidate
293
+ when its score >= (per-class threshold - per-class bonus)."""
294
+ if len(scores) == 0:
295
+ return np.zeros(0, dtype=bool)
296
+ thr = self._conf_thres_array[cls_ids]
297
+ keep = scores >= thr
298
+ for c in np.unique(cls_ids):
299
+ b = float(self._bonus_array[c])
300
+ if b <= 0.0:
301
+ continue
302
+ cm = cls_ids == c
303
+ if keep[cm].any():
304
+ continue
305
+ idx = np.where(cm)[0]
306
+ top = int(idx[int(np.argmax(scores[idx]))])
307
+ if scores[top] >= self._conf_thres_array[c] - b:
308
+ keep[top] = True
309
+ return keep
310
+
311
+ def _per_view_pipeline(self, boxes: np.ndarray, scores: np.ndarray,
312
+ cls_ids: np.ndarray, orig_size: tuple[int, int]
313
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
314
+ boxes, scores, cls_ids = self._filter_sane_boxes(
315
+ boxes, scores, cls_ids, orig_size
316
+ )
317
+ if len(boxes) == 0:
318
+ return boxes, scores, cls_ids
319
+ if len(boxes) > 1:
320
+ keep = self._hard_nms(boxes, scores, self.iou_thres)
321
+ boxes, scores, cls_ids = boxes[keep], scores[keep], cls_ids[keep]
322
+ if len(scores) > self.max_det:
323
+ top = np.argsort(-scores)[: self.max_det]
324
+ boxes, scores, cls_ids = boxes[top], scores[top], cls_ids[top]
325
+ if len(boxes) > 1:
326
+ boxes, scores, cls_ids = self._cross_class_dedup_op(
327
+ boxes, scores, cls_ids, self.cross_iou_thresh
328
+ )
329
+ return boxes, scores, cls_ids
330
+
331
+ def _decode_final_dets(self, preds: np.ndarray, ratio: float,
332
+ pad: tuple[float, float],
333
+ orig_size: tuple[int, int]) -> list[BoundingBox]:
334
+ if preds.ndim == 3 and preds.shape[0] == 1:
335
+ preds = preds[0]
336
+ if preds.ndim != 2 or preds.shape[1] < 6:
337
+ raise ValueError(f"Unexpected ONNX final-det output shape: {preds.shape}")
338
+
339
+ boxes = preds[:, :4].astype(np.float32)
340
+ scores = preds[:, 4].astype(np.float32)
341
+ cls_ids = self._to_competition_cls(preds[:, 5].astype(np.int32))
342
+ valid_cls = cls_ids >= 0
343
+ boxes = boxes[valid_cls]
344
+ scores = scores[valid_cls]
345
+ cls_ids = cls_ids[valid_cls]
346
+
347
+ keep = self._conf_filter_mask(scores, cls_ids)
348
+ boxes = boxes[keep]
349
+ scores = scores[keep]
350
+ cls_ids = cls_ids[keep]
351
+ if len(boxes) == 0:
352
+ return []
353
+
354
+ pad_w, pad_h = pad
355
+ boxes[:, [0, 2]] -= pad_w
356
+ boxes[:, [1, 3]] -= pad_h
357
+ boxes /= ratio
358
+ boxes = self._clip_boxes(boxes, orig_size)
359
+
360
+ boxes, scores, cls_ids = self._per_view_pipeline(
361
+ boxes, scores, cls_ids, orig_size
362
+ )
363
+ return self._build_results(boxes, scores, cls_ids)
364
+
365
+ def _decode_raw_yolo(self, preds: np.ndarray, ratio: float,
366
+ pad: tuple[float, float],
367
+ orig_size: tuple[int, int]) -> list[BoundingBox]:
368
+ if preds.ndim != 3 or preds.shape[0] != 1:
369
+ raise ValueError(f"Unexpected raw ONNX output shape: {preds.shape}")
370
+ preds = preds[0]
371
+ if preds.shape[0] <= 16 and preds.shape[1] > preds.shape[0]:
372
+ preds = preds.T
373
+ if preds.ndim != 2 or preds.shape[1] < 5:
374
+ raise ValueError(f"Unexpected raw output shape: {preds.shape}")
375
+
376
+ boxes_xywh = preds[:, :4].astype(np.float32)
377
+ cls_part = preds[:, 4:].astype(np.float32)
378
+ if cls_part.shape[1] == 1:
379
+ scores = cls_part[:, 0]
380
+ cls_ids = np.zeros(len(scores), dtype=np.int32)
381
+ else:
382
+ cls_ids = np.argmax(cls_part, axis=1).astype(np.int32)
383
+ scores = cls_part[np.arange(len(cls_part)), cls_ids]
384
+ cls_ids = self._to_competition_cls(cls_ids)
385
+ valid_cls = cls_ids >= 0
386
+ boxes_xywh = boxes_xywh[valid_cls]
387
+ scores = scores[valid_cls]
388
+ cls_ids = cls_ids[valid_cls]
389
+
390
+ keep = self._conf_filter_mask(scores, cls_ids)
391
+ boxes_xywh = boxes_xywh[keep]
392
+ scores = scores[keep]
393
+ cls_ids = cls_ids[keep]
394
+ if len(boxes_xywh) == 0:
395
+ return []
396
+ boxes = self._xywh_to_xyxy(boxes_xywh)
397
+
398
+ pad_w, pad_h = pad
399
+ boxes[:, [0, 2]] -= pad_w
400
+ boxes[:, [1, 3]] -= pad_h
401
+ boxes /= ratio
402
+ boxes = self._clip_boxes(boxes, orig_size)
403
+
404
+ boxes, scores, cls_ids = self._per_view_pipeline(
405
+ boxes, scores, cls_ids, orig_size
406
+ )
407
+ return self._build_results(boxes, scores, cls_ids)
408
+
409
+ @staticmethod
410
+ def _build_results(boxes: np.ndarray, scores: np.ndarray,
411
+ cls_ids: np.ndarray) -> list[BoundingBox]:
412
+ results: list[BoundingBox] = []
413
+ for box, conf, cls_id in zip(boxes, scores, cls_ids):
414
+ x1, y1, x2, y2 = box.tolist()
415
+ if x2 <= x1 or y2 <= y1:
416
+ continue
417
+ results.append(
418
+ BoundingBox(
419
+ x1=int(math.floor(x1)),
420
+ y1=int(math.floor(y1)),
421
+ x2=int(math.ceil(x2)),
422
+ y2=int(math.ceil(y2)),
423
+ cls_id=int(cls_id),
424
+ conf=float(conf),
425
+ )
426
+ )
427
+ return results
428
+
429
+ def _postprocess(self, output: np.ndarray, ratio: float,
430
+ pad: tuple[float, float],
431
+ orig_size: tuple[int, int]) -> list[BoundingBox]:
432
+ if output.ndim == 2 and output.shape[1] >= 6:
433
+ return self._decode_final_dets(output, ratio, pad, orig_size)
434
+ if output.ndim == 3 and output.shape[0] == 1 and output.shape[2] == 6:
435
+ return self._decode_final_dets(output, ratio, pad, orig_size)
436
+ return self._decode_raw_yolo(output, ratio, pad, orig_size)
437
+
438
+ def _predict_single(self, image: np.ndarray) -> list[BoundingBox]:
439
+ if image is None:
440
+ raise ValueError("Input image is None")
441
+ if not isinstance(image, np.ndarray):
442
+ raise TypeError(f"Input is not numpy array: {type(image)}")
443
+ if image.ndim != 3:
444
+ raise ValueError(f"Expected HWC image, got shape={image.shape}")
445
+ if image.shape[2] != 3:
446
+ raise ValueError(f"Expected 3 channels, got shape={image.shape}")
447
+ if image.dtype != np.uint8:
448
+ image = image.astype(np.uint8)
449
+
450
+ input_tensor, ratio, pad, orig_size = self._preprocess(image)
451
+ expected = (1, 3, self.input_height, self.input_width)
452
+ if input_tensor.shape != expected:
453
+ raise ValueError(
454
+ f"Bad input tensor shape={input_tensor.shape}, expected={expected}"
455
+ )
456
+
457
+ outputs = self.session.run(self.output_names, {self.input_name: input_tensor})
458
+ return self._postprocess(outputs[0], ratio, pad, orig_size)
459
+
460
+ def _predict_tta(self, image: np.ndarray) -> list[BoundingBox]:
461
+ boxes_orig = self._predict_single(image)
462
+ flipped = cv2.flip(image, 1)
463
+ boxes_flip = self._predict_single(flipped)
464
+ w = image.shape[1]
465
+ boxes_flip = [
466
+ BoundingBox(
467
+ x1=w - b.x2, y1=b.y1, x2=w - b.x1, y2=b.y2,
468
+ cls_id=b.cls_id, conf=b.conf,
469
+ )
470
+ for b in boxes_flip
471
+ ]
472
+ all_boxes = boxes_orig + boxes_flip
473
+ if not all_boxes:
474
+ return []
475
+
476
+ coords = np.array(
477
+ [[b.x1, b.y1, b.x2, b.y2] for b in all_boxes], dtype=np.float32
478
+ )
479
+ scores = np.array([b.conf for b in all_boxes], dtype=np.float32)
480
+ cls_ids = np.array([b.cls_id for b in all_boxes], dtype=np.int32)
481
+
482
+ hard_keep = self._per_class_hard_nms(coords, scores, cls_ids, self.iou_thres)
483
+ if len(hard_keep) == 0:
484
+ return []
485
+ if len(hard_keep) > self.max_det:
486
+ top = np.argsort(-scores[hard_keep])[: self.max_det]
487
+ hard_keep = hard_keep[top]
488
+ boosted = self._max_score_per_cluster(
489
+ coords[hard_keep], cls_ids[hard_keep],
490
+ coords, scores, cls_ids, self.iou_thres,
491
+ )
492
+
493
+ kept_coords = coords[hard_keep]
494
+ kept_cls = cls_ids[hard_keep]
495
+ if len(kept_coords) > 1:
496
+ kept_coords, boosted, kept_cls = self._cross_class_dedup_op(
497
+ kept_coords, boosted, kept_cls, self.cross_iou_thresh
498
+ )
499
+
500
+ return [
501
+ BoundingBox(
502
+ x1=int(math.floor(kept_coords[j, 0])),
503
+ y1=int(math.floor(kept_coords[j, 1])),
504
+ x2=int(math.ceil(kept_coords[j, 2])),
505
+ y2=int(math.ceil(kept_coords[j, 3])),
506
+ cls_id=int(kept_cls[j]),
507
+ conf=float(boosted[j]),
508
+ )
509
+ for j in range(len(kept_coords))
510
+ ]
511
+
512
+ def predict_batch(self, batch_images: list[ndarray], offset: int,
513
+ n_keypoints: int) -> list[TVFrameResult]:
514
+ results: list[TVFrameResult] = []
515
+ for frame_number_in_batch, image in enumerate(batch_images):
516
+ try:
517
+ boxes = self._predict_tta(image)
518
+ except Exception as e:
519
+ print(f"Inference failed for frame {offset + frame_number_in_batch}: {e}")
520
+ boxes = []
521
+ results.append(
522
+ TVFrameResult(
523
+ frame_id=offset + frame_number_in_batch,
524
+ boxes=boxes,
525
+ keypoints=[(0, 0) for _ in range(max(0, int(n_keypoints)))],
526
+ )
527
+ )
528
+ return results
weights.onnx ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:94c540f949e7f8a714bd44b26c86d152de342827d7ce030541742907093aac42
3
+ size 19358729