gtensorapp commited on
Commit
d9a73d7
·
verified ·
1 Parent(s): 8e69aa2

Detect-fire miner

Browse files
Files changed (3) hide show
  1. chute_config.yml +23 -0
  2. miner.py +790 -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
+ # CPU onnxruntime pinned to the conformity checker's version (Dockerfile.latency_conformity):
6
+ # identical kernels => identical boxes between the live chute and the offline FAIL_OUTPUT check.
7
+ - pip install 'numpy>=1.23,<3' 'onnxruntime==1.26.0' 'opencv-python-headless>=4.8' 'pillow>=9.5' 'huggingface_hub>=0.24' 'pydantic>=2.0' 'pyyaml>=6.0' 'aiohttp>=3.9'
8
+ - pip install torch==2.8.0 torchvision==0.23.0 --index-url https://download.pytorch.org/whl/cu128
9
+
10
+ NodeSelector:
11
+ gpu_count: 1
12
+ min_vram_gb_per_gpu: 16
13
+ max_hourly_price_per_gpu: 2
14
+ include:
15
+ - pro_6000
16
+
17
+ Chute:
18
+ timeout_seconds: 900
19
+ concurrency: 4
20
+ max_instances: 5
21
+ scaling_threshold: 0.5
22
+ shutdown_after_seconds: 288000
23
+ tee: true
miner.py ADDED
@@ -0,0 +1,790 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from pathlib import Path
2
+ import math
3
+ import cv2
4
+ import numpy as np
5
+ import onnxruntime as ort
6
+ from numpy import ndarray
7
+ from pydantic import BaseModel
8
+
9
+ class BoundingBox(BaseModel):
10
+ x1: int
11
+ y1: int
12
+ x2: int
13
+ y2: int
14
+ cls_id: int
15
+ conf: float
16
+
17
+ class TVFrameResult(BaseModel):
18
+ frame_id: int
19
+ boxes: list[BoundingBox]
20
+ keypoints: list[tuple[int, int]]
21
+
22
+ class Miner:
23
+ class_names = ['fire', 'smoke', 'fire extinguisher']
24
+ _model_class_order = ['fire', 'smoke', 'fire extinguisher']
25
+ iou_thres = 0.2
26
+ cross_iou_thresh = 1.01
27
+ max_det = 300
28
+ _conf_thres_array = np.array([0.14, 0.16, 0.1], dtype=np.float32)
29
+ _bonus_array = np.array([0.06, 0.02, 0], dtype=np.float32)
30
+ _bonus_classes = tuple((int(c) for c in np.nonzero(_bonus_array > 0.0)[0]))
31
+ min_box_area = 100
32
+ min_side = 8
33
+ max_aspect_ratio = 8
34
+ smoke_merge_overlap = 0.6
35
+ fire_merge_overlap = 0.5
36
+ fire_suppress_overlap = 0.9
37
+ smoke_raw_refine_overlap = 1.9
38
+ smoke_ext_shrink = 0.98
39
+ fire_expand = 1
40
+ use_edge_filter = True
41
+ edge_filter_max_conf = 0.3
42
+ edge_tol = 8
43
+ fire_color_filter_max_conf = 0.0
44
+ fire_ext_color_filter_max_conf = 0.0
45
+ color_filter_min_saturation = 0.06
46
+ use_tta_view_filter = False
47
+ tta_view_filter_max_conf = 0.0
48
+ tta_view_iou_thresh = 0.5
49
+
50
+ def __init__(self, path_hf_repo: Path) -> None:
51
+ model_path = path_hf_repo / 'weights.onnx'
52
+ print('ORT version:', ort.__version__)
53
+ try:
54
+ ort.preload_dlls()
55
+ print('✅ onnxruntime.preload_dlls() success')
56
+ except Exception as e:
57
+ print(f'⚠️ preload_dlls failed: {e}')
58
+ print('ORT available providers BEFORE session:', ort.get_available_providers())
59
+ sess_options = ort.SessionOptions()
60
+ sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
61
+ sess_options.intra_op_num_threads = 2
62
+ sess_options.inter_op_num_threads = 1
63
+ sess_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
64
+ try:
65
+ self.session = ort.InferenceSession(str(model_path), sess_options=sess_options, providers=["CPUExecutionProvider"])
66
+ print('✅ Created ORT session with preferred CUDA provider list')
67
+ except Exception as e:
68
+ print(f'⚠️ CUDA session creation failed, falling back to CPU: {e}')
69
+ self.session = ort.InferenceSession(str(model_path), sess_options=sess_options, providers=['CPUExecutionProvider'])
70
+ print('ORT session providers:', self.session.get_providers())
71
+ model_class_order = self._read_model_class_order()
72
+ if model_class_order is None:
73
+ model_class_order = list(self._model_class_order)
74
+ print(f'cls order: no usable ONNX metadata, FALLBACK {model_class_order}')
75
+ else:
76
+ print(f'cls order: from ONNX metadata {model_class_order}')
77
+ self.cls_remap = np.array([self.class_names.index(n) for n in model_class_order], dtype=np.int32)
78
+ for inp in self.session.get_inputs():
79
+ print('INPUT:', inp.name, inp.shape, inp.type)
80
+ for out in self.session.get_outputs():
81
+ print('OUTPUT:', out.name, out.shape, out.type)
82
+ self.input_name = self.session.get_inputs()[0].name
83
+ self.output_names = [output.name for output in self.session.get_outputs()]
84
+ self.input_shape = self.session.get_inputs()[0].shape
85
+ self.input_height = self._safe_dim(self.input_shape[2], default=1280)
86
+ self.input_width = self._safe_dim(self.input_shape[3], default=1280)
87
+ self._blob = np.empty((1, 3, self.input_height, self.input_width), dtype=np.float32)
88
+ self.use_tta = False
89
+ print(f'✅ ONNX model loaded from: {model_path}')
90
+ print(f'✅ ONNX providers: {self.session.get_providers()}')
91
+ print(f'✅ ONNX input: name={self.input_name}, shape={self.input_shape}')
92
+ print('per-class conf: ' + ', '.join((f'{n}={t:.3f}' for n, t in zip(self.class_names, self._conf_thres_array.tolist()))))
93
+ self._warmup()
94
+
95
+ def _warmup(self, iters: int=3) -> None:
96
+ try:
97
+ dummy = np.zeros((720, 1280, 3), dtype=np.uint8)
98
+ for _ in range(max(1, iters)):
99
+ self.predict_batch(batch_images=[dummy], offset=0, n_keypoints=0)
100
+ print(f'✅ warmup: {iters} dummy predict_batch call(s) done')
101
+ except Exception as e:
102
+ print(f'⚠️ warmup skipped: {e}')
103
+
104
+ def __repr__(self) -> str:
105
+ return f'ONNXRuntime(session={type(self.session).__name__}, providers={self.session.get_providers()})'
106
+
107
+ @staticmethod
108
+ def _safe_dim(value, default: int) -> int:
109
+ return value if isinstance(value, int) and value > 0 else default
110
+
111
+ def _read_model_class_order(self) -> list[str] | None:
112
+ try:
113
+ import ast
114
+ meta = self.session.get_modelmeta().custom_metadata_map
115
+ names = ast.literal_eval(meta['names'])
116
+ if isinstance(names, dict):
117
+ order = [str(names[i]) for i in sorted(names)]
118
+ else:
119
+ order = [str(n) for n in names]
120
+ except Exception as e:
121
+ print(f'cls order: could not read ONNX names metadata ({e})')
122
+ return None
123
+ if sorted(order) != sorted(self.class_names):
124
+ print(f'cls order: ONNX names {order} do not match expected classes {self.class_names}; ignoring metadata')
125
+ return None
126
+ return order
127
+
128
+ def _letterbox(self, image: ndarray, new_shape: tuple[int, int], color=(114, 114, 114)) -> tuple[ndarray, float, tuple[float, float]]:
129
+ h, w = image.shape[:2]
130
+ new_w, new_h = new_shape
131
+ ratio = min(new_w / w, new_h / h)
132
+ resized_w = int(round(w * ratio))
133
+ resized_h = int(round(h * ratio))
134
+ if (resized_w, resized_h) != (w, h):
135
+ interp = cv2.INTER_CUBIC if ratio > 1.0 else cv2.INTER_LINEAR
136
+ image = cv2.resize(image, (resized_w, resized_h), interpolation=interp)
137
+ dw = (new_w - resized_w) / 2.0
138
+ dh = (new_h - resized_h) / 2.0
139
+ left = int(round(dw - 0.1))
140
+ right = int(round(dw + 0.1))
141
+ top = int(round(dh - 0.1))
142
+ bottom = int(round(dh + 0.1))
143
+ padded = cv2.copyMakeBorder(image, top, bottom, left, right, borderType=cv2.BORDER_CONSTANT, value=color)
144
+ return (padded, ratio, (dw, dh))
145
+
146
+ def _preprocess(self, image: ndarray) -> tuple[np.ndarray, float, tuple[float, float], tuple[int, int]]:
147
+ orig_h, orig_w = image.shape[:2]
148
+ img, ratio, pad = self._letterbox(image, (self.input_width, self.input_height))
149
+ np.multiply(img.transpose(2, 0, 1)[::-1], 1.0 / 255.0, out=self._blob[0])
150
+ return (self._blob, ratio, pad, (orig_w, orig_h))
151
+
152
+ @staticmethod
153
+ def _clip_boxes(boxes: np.ndarray, image_size: tuple[int, int]) -> np.ndarray:
154
+ w, h = image_size
155
+ boxes[:, 0] = np.clip(boxes[:, 0], 0, w - 1)
156
+ boxes[:, 1] = np.clip(boxes[:, 1], 0, h - 1)
157
+ boxes[:, 2] = np.clip(boxes[:, 2], 0, w - 1)
158
+ boxes[:, 3] = np.clip(boxes[:, 3], 0, h - 1)
159
+ return boxes
160
+
161
+ @staticmethod
162
+ def _xywh_to_xyxy(boxes: np.ndarray) -> np.ndarray:
163
+ out = np.empty_like(boxes)
164
+ out[:, 0] = boxes[:, 0] - boxes[:, 2] / 2.0
165
+ out[:, 1] = boxes[:, 1] - boxes[:, 3] / 2.0
166
+ out[:, 2] = boxes[:, 0] + boxes[:, 2] / 2.0
167
+ out[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.0
168
+ return out
169
+
170
+ @staticmethod
171
+ def _hard_nms(boxes: np.ndarray, scores: np.ndarray, iou_thresh: float) -> np.ndarray:
172
+ n = len(boxes)
173
+ if n == 0:
174
+ return np.array([], dtype=np.intp)
175
+ if n == 1:
176
+ return np.zeros(1, dtype=np.intp)
177
+ areas = np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) * np.maximum(0.0, boxes[:, 3] - boxes[:, 1])
178
+ order = np.argsort(-scores)
179
+ keep: list[int] = []
180
+ while len(order) > 0:
181
+ i = int(order[0])
182
+ keep.append(i)
183
+ if len(order) == 1:
184
+ break
185
+ rest = order[1:]
186
+ bx = boxes[rest]
187
+ xx1 = np.maximum(boxes[i, 0], bx[:, 0])
188
+ yy1 = np.maximum(boxes[i, 1], bx[:, 1])
189
+ xx2 = np.minimum(boxes[i, 2], bx[:, 2])
190
+ yy2 = np.minimum(boxes[i, 3], bx[:, 3])
191
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
192
+ a_i = max(0.0, boxes[i, 2] - boxes[i, 0]) * max(0.0, boxes[i, 3] - boxes[i, 1])
193
+ iou = inter / (a_i + areas[rest] - inter + 1e-07)
194
+ order = rest[iou <= iou_thresh]
195
+ return np.array(keep, dtype=np.intp)
196
+
197
+ def _per_class_hard_nms(self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray, iou_thresh: float) -> np.ndarray:
198
+ if len(boxes) == 0:
199
+ return np.array([], dtype=np.intp)
200
+ all_keep: list[int] = []
201
+ for c in np.unique(cls_ids):
202
+ indices = np.where(cls_ids == c)[0]
203
+ if len(indices) == 1:
204
+ all_keep.append(int(indices[0]))
205
+ continue
206
+ keep = self._hard_nms(boxes[indices], scores[indices], iou_thresh)
207
+ all_keep.extend(indices[keep].tolist())
208
+ all_keep.sort()
209
+ return np.array(all_keep, dtype=np.intp)
210
+
211
+ def _cross_class_dedup_op(self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray, iou_thresh: float) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
212
+ n = len(boxes)
213
+ if n <= 1:
214
+ return (boxes, scores, cls_ids)
215
+ boxes = np.asarray(boxes, dtype=np.float32)
216
+ scores = np.asarray(scores, dtype=np.float32)
217
+ cls_ids = np.asarray(cls_ids, dtype=np.int32)
218
+ areas = np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) * np.maximum(0.0, boxes[:, 3] - boxes[:, 1])
219
+ margins = scores - self._conf_thres_array[cls_ids]
220
+ order = np.lexsort((-areas, -margins))
221
+ if iou_thresh >= 1.0:
222
+ return (boxes[order], scores[order], cls_ids[order])
223
+ suppressed = np.zeros(n, dtype=bool)
224
+ keep: list[int] = []
225
+ for i in order:
226
+ if suppressed[i]:
227
+ continue
228
+ keep.append(int(i))
229
+ bi = boxes[i]
230
+ xx1 = np.maximum(bi[0], boxes[:, 0])
231
+ yy1 = np.maximum(bi[1], boxes[:, 1])
232
+ xx2 = np.minimum(bi[2], boxes[:, 2])
233
+ yy2 = np.minimum(bi[3], boxes[:, 3])
234
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
235
+ a_i = max(1e-07, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
236
+ iou = inter / (a_i + areas - inter + 1e-07)
237
+ dup = iou > iou_thresh
238
+ dup[i] = False
239
+ suppressed |= dup
240
+ keep_idx = np.array(keep, dtype=np.intp)
241
+ return (boxes[keep_idx], scores[keep_idx], cls_ids[keep_idx])
242
+
243
+ def _merge_class_boxes(self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray, target_cls: int, overlap: float) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
244
+ if overlap > 1.0:
245
+ return (boxes, scores, cls_ids)
246
+ idx = np.where(cls_ids == target_cls)[0]
247
+ if len(idx) <= 1:
248
+ return (boxes, scores, cls_ids)
249
+ sb = boxes[idx].astype(np.float32).tolist()
250
+ ss = scores[idx].astype(np.float32).tolist()
251
+ merged_any = True
252
+ while merged_any and len(sb) > 1:
253
+ merged_any = False
254
+ for i in range(len(sb)):
255
+ for j in range(i + 1, len(sb)):
256
+ a, b = (sb[i], sb[j])
257
+ ix1 = max(a[0], b[0])
258
+ iy1 = max(a[1], b[1])
259
+ ix2 = min(a[2], b[2])
260
+ iy2 = min(a[3], b[3])
261
+ inter = max(0.0, ix2 - ix1) * max(0.0, iy2 - iy1)
262
+ area_a = max(0.0, a[2] - a[0]) * max(0.0, a[3] - a[1])
263
+ area_b = max(0.0, b[2] - b[0]) * max(0.0, b[3] - b[1])
264
+ smaller = min(area_a, area_b)
265
+ if inter / (smaller + 1e-07) >= overlap:
266
+ sb[i] = [min(a[0], b[0]), min(a[1], b[1]), max(a[2], b[2]), max(a[3], b[3])]
267
+ ss[i] = max(ss[i], ss[j])
268
+ del sb[j]
269
+ del ss[j]
270
+ merged_any = True
271
+ break
272
+ if merged_any:
273
+ break
274
+ other = cls_ids != target_cls
275
+ new_boxes = np.concatenate([boxes[other].astype(np.float32), np.array(sb, dtype=np.float32).reshape(-1, 4)])
276
+ new_scores = np.concatenate([scores[other].astype(np.float32), np.array(ss, dtype=np.float32)])
277
+ new_cls = np.concatenate([cls_ids[other].astype(np.int32), np.full(len(sb), target_cls, dtype=np.int32)])
278
+ return (new_boxes, new_scores, new_cls)
279
+
280
+ def _suppress_contained_lower_conf(self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray, target_cls: int, overlap: float) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
281
+ if overlap > 1.0:
282
+ return (boxes, scores, cls_ids)
283
+ idx = np.where(cls_ids == target_cls)[0]
284
+ if len(idx) <= 1:
285
+ return (boxes, scores, cls_ids)
286
+ order = idx[np.argsort(-scores[idx])]
287
+ remove: set[int] = set()
288
+ for a in range(len(order)):
289
+ i = int(order[a])
290
+ if i in remove:
291
+ continue
292
+ bi = boxes[i]
293
+ area_i = max(1e-07, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
294
+ for b in range(a + 1, len(order)):
295
+ j = int(order[b])
296
+ if j in remove:
297
+ continue
298
+ bj = boxes[j]
299
+ ix1 = max(bi[0], bj[0])
300
+ iy1 = max(bi[1], bj[1])
301
+ ix2 = min(bi[2], bj[2])
302
+ iy2 = min(bi[3], bj[3])
303
+ inter = max(0.0, ix2 - ix1) * max(0.0, iy2 - iy1)
304
+ if inter <= 0.0:
305
+ continue
306
+ area_j = max(1e-07, float((bj[2] - bj[0]) * (bj[3] - bj[1])))
307
+ if inter / (min(area_i, area_j) + 1e-07) >= overlap:
308
+ remove.add(j)
309
+ if not remove:
310
+ return (boxes, scores, cls_ids)
311
+ keep = np.array([k not in remove for k in range(len(boxes))], dtype=bool)
312
+ return (boxes[keep], scores[keep], cls_ids[keep])
313
+
314
+ def _merge_same_class_boxes(self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
315
+ boxes, scores, cls_ids = self._merge_class_boxes(boxes, scores, cls_ids, self.class_names.index('smoke'), self.smoke_merge_overlap)
316
+ boxes, scores, cls_ids = self._merge_class_boxes(boxes, scores, cls_ids, self.class_names.index('fire'), self.fire_merge_overlap)
317
+ boxes, scores, cls_ids = self._suppress_contained_lower_conf(boxes, scores, cls_ids, self.class_names.index('fire'), self.fire_suppress_overlap)
318
+ return (boxes, scores, cls_ids)
319
+
320
+ def _merge_smoke_boxes(self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
321
+ return self._merge_same_class_boxes(boxes, scores, cls_ids)
322
+
323
+ @staticmethod
324
+ def _max_score_per_cluster(post_boxes: np.ndarray, post_cls: np.ndarray, full_boxes: np.ndarray, full_scores: np.ndarray, full_cls: np.ndarray, iou_thresh: float) -> np.ndarray:
325
+ n = len(post_boxes)
326
+ if n == 0:
327
+ return np.empty(0, dtype=np.float32)
328
+ full_areas = np.maximum(0.0, full_boxes[:, 2] - full_boxes[:, 0]) * np.maximum(0.0, full_boxes[:, 3] - full_boxes[:, 1])
329
+ out = np.empty(n, dtype=np.float32)
330
+ for i in range(n):
331
+ bi = post_boxes[i]
332
+ xx1 = np.maximum(bi[0], full_boxes[:, 0])
333
+ yy1 = np.maximum(bi[1], full_boxes[:, 1])
334
+ xx2 = np.minimum(bi[2], full_boxes[:, 2])
335
+ yy2 = np.minimum(bi[3], full_boxes[:, 3])
336
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
337
+ a_i = max(0.0, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
338
+ iou = inter / (a_i + full_areas - inter + 1e-07)
339
+ cluster = (iou >= iou_thresh) & (full_cls == post_cls[i])
340
+ out[i] = float(np.max(full_scores[cluster])) if np.any(cluster) else 0.0
341
+ return out
342
+
343
+ def _conf_filter_mask(self, scores: np.ndarray, cls_ids: np.ndarray) -> np.ndarray:
344
+ if len(scores) == 0:
345
+ return np.zeros(0, dtype=bool)
346
+ thr = self._conf_thres_array[cls_ids]
347
+ keep = scores >= thr
348
+ for c in self._bonus_classes:
349
+ b = float(self._bonus_array[c])
350
+ idx = np.where(cls_ids == c)[0]
351
+ if len(idx) == 0:
352
+ continue
353
+ if keep[idx].any():
354
+ continue
355
+ top = int(idx[int(np.argmax(scores[idx]))])
356
+ if scores[top] >= self._conf_thres_array[c] - b:
357
+ keep[top] = True
358
+ return keep
359
+
360
+ def _filter_sane_boxes(self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray, orig_size: tuple[int, int]) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
361
+ if len(boxes) == 0:
362
+ return (boxes, scores, cls_ids)
363
+ orig_w, orig_h = orig_size
364
+ image_area = float(orig_w * orig_h)
365
+ keep = []
366
+ for i, box in enumerate(boxes):
367
+ x1, y1, x2, y2 = box.tolist()
368
+ bw = x2 - x1
369
+ bh = y2 - y1
370
+ if bw <= 0 or bh <= 0:
371
+ continue
372
+ if bw < self.min_side or bh < self.min_side:
373
+ continue
374
+ area = bw * bh
375
+ if area < self.min_box_area:
376
+ continue
377
+ if area > 0.95 * image_area:
378
+ continue
379
+ ar = max(bw / max(bh, 1e-06), bh / max(bw, 1e-06))
380
+ if ar > self.max_aspect_ratio:
381
+ continue
382
+ keep.append(i)
383
+ if not keep:
384
+ return (np.empty((0, 4), dtype=np.float32), np.empty((0,), dtype=np.float32), np.empty((0,), dtype=np.int32))
385
+ k = np.array(keep, dtype=np.intp)
386
+ return (boxes[k], scores[k], cls_ids[k])
387
+
388
+ def _per_view_pipeline(self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
389
+ if len(boxes) > 1:
390
+ keep = self._per_class_hard_nms(boxes, scores, cls_ids, self.iou_thres)
391
+ boxes, scores, cls_ids = (boxes[keep], scores[keep], cls_ids[keep])
392
+ if len(scores) > self.max_det:
393
+ top = np.argsort(-scores)[:self.max_det]
394
+ boxes, scores, cls_ids = (boxes[top], scores[top], cls_ids[top])
395
+ if len(boxes) > 1:
396
+ boxes, scores, cls_ids = self._cross_class_dedup_op(boxes, scores, cls_ids, self.cross_iou_thresh)
397
+ if len(boxes) > 1:
398
+ boxes, scores, cls_ids = self._merge_same_class_boxes(boxes, scores, cls_ids)
399
+ return (boxes, scores, cls_ids)
400
+
401
+ @staticmethod
402
+ def _roi_for_box(image: np.ndarray, box: BoundingBox) -> np.ndarray | None:
403
+ h, w = image.shape[:2]
404
+ x1 = max(0, int(math.floor(box.x1)))
405
+ y1 = max(0, int(math.floor(box.y1)))
406
+ x2 = min(w, int(math.ceil(box.x2)))
407
+ y2 = min(h, int(math.ceil(box.y2)))
408
+ if x2 <= x1 or y2 <= y1:
409
+ return None
410
+ roi = image[y1:y2, x1:x2]
411
+ return roi if roi.size else None
412
+
413
+ def _roi_is_near_grayscale(self, roi: np.ndarray) -> bool:
414
+ mx = roi.max(axis=2).astype(np.float32)
415
+ mn = roi.min(axis=2).astype(np.float32)
416
+ sat = (mx - mn) / (mx + 1e-06)
417
+ return float(sat.mean()) < self.color_filter_min_saturation
418
+
419
+ @staticmethod
420
+ def _passes_fire_color(roi: np.ndarray) -> bool:
421
+ blue = roi[:, :, 0].astype(np.float32)
422
+ green = roi[:, :, 1].astype(np.float32)
423
+ red = roi[:, :, 2].astype(np.float32)
424
+ mean_r = float(np.mean(red))
425
+ max_rgb = float(max(np.max(red), np.max(green), np.max(blue)))
426
+ bright_frac = float(np.mean(np.max(roi, axis=2) >= 150))
427
+ if max_rgb >= 200.0 and bright_frac >= 0.01:
428
+ return True
429
+ warm = (red > green + 10.0) & (red > blue + 10.0)
430
+ warm_frac = float(np.mean(warm))
431
+ r_minus_g = mean_r - float(np.mean(green))
432
+ if warm_frac >= 0.05 and (max_rgb >= 120.0 or mean_r >= 120.0 or warm_frac >= 0.15):
433
+ return True
434
+ if bright_frac >= 0.12 and r_minus_g >= 2.0:
435
+ return True
436
+ return False
437
+
438
+ @staticmethod
439
+ def _passes_fire_ext_red_color(roi: np.ndarray) -> bool:
440
+ blue = roi[:, :, 0].astype(np.float32)
441
+ green = roi[:, :, 1].astype(np.float32)
442
+ red = roi[:, :, 2].astype(np.float32)
443
+ red_dom = float(np.mean((red > green + 10.0) & (red > blue + 10.0)))
444
+ if red_dom >= 0.03:
445
+ return True
446
+ if float(np.mean(red)) - float(np.mean(green)) >= 0.0 and float(np.mean(red)) >= 50.0:
447
+ return True
448
+ return False
449
+
450
+ def _remove_edge_low_conf(self, results: list[BoundingBox], orig_size: tuple[int, int]) -> list[BoundingBox]:
451
+ if not self.use_edge_filter or self.edge_filter_max_conf <= 0.0 or (not results):
452
+ return results
453
+ w, h = orig_size
454
+ tol = self.edge_tol
455
+ out: list[BoundingBox] = []
456
+ for b in results:
457
+ on_edge = b.x1 <= tol or b.y1 <= tol or b.x2 >= w - 1 - tol or (b.y2 >= h - 1 - tol)
458
+ if on_edge and b.conf <= self.edge_filter_max_conf:
459
+ continue
460
+ out.append(b)
461
+ return out
462
+
463
+ def _views_corroborated(self, post_boxes: np.ndarray, post_cls: np.ndarray, full_boxes: np.ndarray, full_cls: np.ndarray, full_views: np.ndarray, iou_thresh: float) -> np.ndarray:
464
+ n = len(post_boxes)
465
+ if n == 0:
466
+ return np.zeros(0, dtype=bool)
467
+ full_areas = np.maximum(0.0, full_boxes[:, 2] - full_boxes[:, 0]) * np.maximum(0.0, full_boxes[:, 3] - full_boxes[:, 1])
468
+ out = np.zeros(n, dtype=bool)
469
+ for i in range(n):
470
+ bi = post_boxes[i]
471
+ xx1 = np.maximum(bi[0], full_boxes[:, 0])
472
+ yy1 = np.maximum(bi[1], full_boxes[:, 1])
473
+ xx2 = np.minimum(bi[2], full_boxes[:, 2])
474
+ yy2 = np.minimum(bi[3], full_boxes[:, 3])
475
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
476
+ a_i = max(0.0, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
477
+ iou = inter / (a_i + full_areas - inter + 1e-07)
478
+ mask = (iou >= iou_thresh) & (full_cls == post_cls[i])
479
+ if np.any(mask):
480
+ out[i] = len(np.unique(full_views[mask])) >= 2
481
+ return out
482
+
483
+ def _filter_low_conf_by_color(self, image: np.ndarray, results: list[BoundingBox]) -> list[BoundingBox]:
484
+ if not results:
485
+ return results
486
+ cls_fire = self.class_names.index('fire')
487
+ cls_ext = self.class_names.index('fire extinguisher')
488
+ out: list[BoundingBox] = []
489
+ for box in results:
490
+ check_fire = box.cls_id == cls_fire and box.conf <= self.fire_color_filter_max_conf
491
+ check_ext = box.cls_id == cls_ext and box.conf <= self.fire_ext_color_filter_max_conf
492
+ if not check_fire and (not check_ext):
493
+ out.append(box)
494
+ continue
495
+ roi = self._roi_for_box(image, box)
496
+ if roi is None or self._roi_is_near_grayscale(roi):
497
+ out.append(box)
498
+ continue
499
+ if check_fire and (not self._passes_fire_color(roi)):
500
+ continue
501
+ if check_ext and (not self._passes_fire_ext_red_color(roi)):
502
+ continue
503
+ out.append(box)
504
+ return out
505
+
506
+ @staticmethod
507
+ def _build_results(boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray) -> list[BoundingBox]:
508
+ results: list[BoundingBox] = []
509
+ for box, conf, cls_id in zip(boxes, scores, cls_ids):
510
+ x1, y1, x2, y2 = box.tolist()
511
+ if x2 <= x1 or y2 <= y1:
512
+ continue
513
+ results.append(BoundingBox(x1=int(math.floor(x1)), y1=int(math.floor(y1)), x2=int(math.ceil(x2)), y2=int(math.ceil(y2)), cls_id=int(cls_id), conf=float(conf)))
514
+ return results
515
+
516
+ @staticmethod
517
+ def _empty_raw() -> tuple[np.ndarray, np.ndarray, np.ndarray]:
518
+ return (np.empty((0, 4), dtype=np.float32), np.empty((0,), dtype=np.float32), np.empty((0,), dtype=np.int32))
519
+
520
+ @staticmethod
521
+ def _iomin(a: np.ndarray, b: np.ndarray) -> float:
522
+ ix1 = max(float(a[0]), float(b[0]))
523
+ iy1 = max(float(a[1]), float(b[1]))
524
+ ix2 = min(float(a[2]), float(b[2]))
525
+ iy2 = min(float(a[3]), float(b[3]))
526
+ inter = max(0.0, ix2 - ix1) * max(0.0, iy2 - iy1)
527
+ area_a = max(0.0, float(a[2] - a[0]) * float(a[3] - a[1]))
528
+ area_b = max(0.0, float(b[2] - b[0]) * float(b[3] - b[1]))
529
+ smaller = min(area_a, area_b)
530
+ return inter / (smaller + 1e-07)
531
+
532
+ def _refine_smoke_from_raw(self, finals: list[BoundingBox], raw_boxes: np.ndarray, raw_scores: np.ndarray, raw_cls: np.ndarray) -> list[BoundingBox]:
533
+ del raw_scores
534
+ if not finals or len(raw_boxes) == 0:
535
+ return finals
536
+ smoke_id = self.class_names.index('smoke')
537
+ raw_smoke = raw_cls == smoke_id
538
+ if not np.any(raw_smoke):
539
+ return finals
540
+ thr = float(self.smoke_raw_refine_overlap)
541
+ if thr > 1.0:
542
+ return self._nms_final_boxes(list(finals))
543
+ cand_boxes = raw_boxes[raw_smoke]
544
+ out: list[BoundingBox] = []
545
+ for b in finals:
546
+ if b.cls_id != smoke_id:
547
+ out.append(b)
548
+ continue
549
+ final_xyxy = np.array([b.x1, b.y1, b.x2, b.y2], dtype=np.float32)
550
+ best_idx = -1
551
+ best_area = None
552
+ for i, rb in enumerate(cand_boxes):
553
+ if self._iomin(final_xyxy, rb) < thr:
554
+ continue
555
+ area = max(0.0, float(rb[2] - rb[0]) * float(rb[3] - rb[1]))
556
+ if best_area is None or area < best_area:
557
+ best_area = area
558
+ best_idx = i
559
+ if best_idx < 0:
560
+ out.append(b)
561
+ continue
562
+ rb = cand_boxes[best_idx]
563
+ out.append(BoundingBox(x1=int(math.floor(rb[0])), y1=int(math.floor(rb[1])), x2=int(math.ceil(rb[2])), y2=int(math.ceil(rb[3])), cls_id=b.cls_id, conf=b.conf))
564
+ return self._nms_final_boxes(out)
565
+
566
+ def _nms_final_boxes(self, out: list[BoundingBox]) -> list[BoundingBox]:
567
+ if len(out) <= 1:
568
+ return out
569
+ boxes = np.array([[b.x1, b.y1, b.x2, b.y2] for b in out], dtype=np.float32)
570
+ scores = np.array([b.conf for b in out], dtype=np.float32)
571
+ cls_ids = np.array([b.cls_id for b in out], dtype=np.int32)
572
+ keep = self._per_class_hard_nms(boxes, scores, cls_ids, self.iou_thres)
573
+ return [out[int(i)] for i in keep]
574
+
575
+ def _rescale_class_boxes(self, finals: list[BoundingBox], orig_size: tuple[int, int]) -> list[BoundingBox]:
576
+ if not finals:
577
+ return finals
578
+ img_w, img_h = orig_size
579
+ fire_id = self.class_names.index('fire')
580
+ smoke_id = self.class_names.index('smoke')
581
+ ext_id = self.class_names.index('fire extinguisher')
582
+ out: list[BoundingBox] = []
583
+ for b in finals:
584
+ x1, y1, x2, y2 = (float(b.x1), float(b.y1), float(b.x2), float(b.y2))
585
+ w = max(0.0, x2 - x1)
586
+ h = max(0.0, y2 - y1)
587
+ if w <= 0.0 or h <= 0.0:
588
+ continue
589
+ if b.cls_id in (smoke_id, ext_id):
590
+ scale = float(self.smoke_ext_shrink)
591
+ nw, nh = (w * scale, h * scale)
592
+ cx = 0.5 * (x1 + x2)
593
+ nx1 = cx - 0.5 * nw
594
+ nx2 = cx + 0.5 * nw
595
+ ny2 = y2
596
+ ny1 = ny2 - nh
597
+ elif b.cls_id == fire_id:
598
+ scale = float(self.fire_expand)
599
+ nw, nh = (w * scale, h * scale)
600
+ cx = 0.5 * (x1 + x2)
601
+ cy = 0.5 * (y1 + y2)
602
+ nx1 = cx - 0.5 * nw
603
+ nx2 = cx + 0.5 * nw
604
+ ny1 = cy - 0.5 * nh
605
+ ny2 = cy + 0.5 * nh
606
+ else:
607
+ out.append(b)
608
+ continue
609
+ nx1 = max(0.0, min(float(img_w), nx1))
610
+ nx2 = max(0.0, min(float(img_w), nx2))
611
+ ny1 = max(0.0, min(float(img_h), ny1))
612
+ ny2 = max(0.0, min(float(img_h), ny2))
613
+ if nx2 <= nx1 or ny2 <= ny1:
614
+ continue
615
+ out.append(BoundingBox(x1=int(math.floor(nx1)), y1=int(math.floor(ny1)), x2=int(math.ceil(nx2)), y2=int(math.ceil(ny2)), cls_id=b.cls_id, conf=b.conf))
616
+ return out
617
+
618
+ def _apply_extra_post(self, finals: list[BoundingBox], raw_boxes: np.ndarray, raw_scores: np.ndarray, raw_cls: np.ndarray, orig_size: tuple[int, int]) -> list[BoundingBox]:
619
+ finals = self._refine_smoke_from_raw(finals, raw_boxes, raw_scores, raw_cls)
620
+ return self._rescale_class_boxes(finals, orig_size)
621
+
622
+ def _candidates_final_dets(self, preds: np.ndarray, ratio: float, pad: tuple[float, float], orig_size: tuple[int, int]) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
623
+ if preds.ndim == 3 and preds.shape[0] == 1:
624
+ preds = preds[0]
625
+ if preds.ndim != 2 or preds.shape[1] < 6:
626
+ raise ValueError(f'Unexpected ONNX final-det output shape: {preds.shape}')
627
+ boxes = preds[:, :4].astype(np.float32)
628
+ scores = preds[:, 4].astype(np.float32)
629
+ cls_ids = preds[:, 5].astype(np.int32)
630
+ cls_ids = self.cls_remap[cls_ids]
631
+ keep = self._conf_filter_mask(scores, cls_ids)
632
+ boxes = boxes[keep]
633
+ scores = scores[keep]
634
+ cls_ids = cls_ids[keep]
635
+ if len(boxes) == 0:
636
+ return self._empty_raw()
637
+ pad_w, pad_h = pad
638
+ boxes[:, [0, 2]] -= pad_w
639
+ boxes[:, [1, 3]] -= pad_h
640
+ boxes /= ratio
641
+ boxes = self._clip_boxes(boxes, orig_size)
642
+ return (boxes, scores, cls_ids)
643
+
644
+ def _candidates_raw_yolo(self, preds: np.ndarray, ratio: float, pad: tuple[float, float], orig_size: tuple[int, int]) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
645
+ if preds.ndim != 3 or preds.shape[0] != 1:
646
+ raise ValueError(f'Unexpected raw ONNX output shape: {preds.shape}')
647
+ preds = preds[0]
648
+ if preds.shape[0] <= 16 and preds.shape[1] > preds.shape[0]:
649
+ preds = preds.T
650
+ if preds.ndim != 2 or preds.shape[1] < 5:
651
+ raise ValueError(f'Unexpected raw output shape: {preds.shape}')
652
+ boxes_xywh = preds[:, :4].astype(np.float32)
653
+ cls_part = preds[:, 4:].astype(np.float32)
654
+ if cls_part.shape[1] == 1:
655
+ scores = cls_part[:, 0]
656
+ cls_ids = np.zeros(len(scores), dtype=np.int32)
657
+ else:
658
+ cls_ids = np.argmax(cls_part, axis=1).astype(np.int32)
659
+ scores = cls_part[np.arange(len(cls_part)), cls_ids]
660
+ cls_ids = self.cls_remap[cls_ids]
661
+ keep = self._conf_filter_mask(scores, cls_ids)
662
+ boxes_xywh = boxes_xywh[keep]
663
+ scores = scores[keep]
664
+ cls_ids = cls_ids[keep]
665
+ if len(boxes_xywh) == 0:
666
+ return self._empty_raw()
667
+ boxes = self._xywh_to_xyxy(boxes_xywh)
668
+ pad_w, pad_h = pad
669
+ boxes[:, [0, 2]] -= pad_w
670
+ boxes[:, [1, 3]] -= pad_h
671
+ boxes /= ratio
672
+ boxes = self._clip_boxes(boxes, orig_size)
673
+ return (boxes, scores, cls_ids)
674
+
675
+ def _pipeline_from_candidates(self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray, orig_size: tuple[int, int]) -> list[BoundingBox]:
676
+ if len(boxes) == 0:
677
+ return []
678
+ boxes, scores, cls_ids = self._filter_sane_boxes(boxes, scores, cls_ids, orig_size)
679
+ if len(boxes) == 0:
680
+ return []
681
+ boxes, scores, cls_ids = self._per_view_pipeline(boxes, scores, cls_ids)
682
+ return self._build_results(boxes, scores, cls_ids)
683
+
684
+ def _decode_final_dets(self, preds: np.ndarray, ratio: float, pad: tuple[float, float], orig_size: tuple[int, int]) -> tuple[list[BoundingBox], tuple[np.ndarray, np.ndarray, np.ndarray]]:
685
+ raw = self._candidates_final_dets(preds, ratio, pad, orig_size)
686
+ return (self._pipeline_from_candidates(*raw, orig_size), raw)
687
+
688
+ def _decode_raw_yolo(self, preds: np.ndarray, ratio: float, pad: tuple[float, float], orig_size: tuple[int, int]) -> tuple[list[BoundingBox], tuple[np.ndarray, np.ndarray, np.ndarray]]:
689
+ raw = self._candidates_raw_yolo(preds, ratio, pad, orig_size)
690
+ return (self._pipeline_from_candidates(*raw, orig_size), raw)
691
+
692
+ def _postprocess(self, output: np.ndarray, ratio: float, pad: tuple[float, float], orig_size: tuple[int, int]) -> tuple[list[BoundingBox], tuple[np.ndarray, np.ndarray, np.ndarray]]:
693
+ if output.ndim == 2 and output.shape[1] >= 6:
694
+ return self._decode_final_dets(output, ratio, pad, orig_size)
695
+ if output.ndim == 3 and output.shape[0] == 1 and (output.shape[2] == 6):
696
+ return self._decode_final_dets(output, ratio, pad, orig_size)
697
+ return self._decode_raw_yolo(output, ratio, pad, orig_size)
698
+
699
+ def _predict_single(self, image: np.ndarray) -> tuple[list[BoundingBox], tuple[np.ndarray, np.ndarray, np.ndarray]]:
700
+ if image is None:
701
+ raise ValueError('Input image is None')
702
+ if not isinstance(image, np.ndarray):
703
+ raise TypeError(f'Input is not numpy array: {type(image)}')
704
+ if image.ndim != 3:
705
+ raise ValueError(f'Expected HWC image, got shape={image.shape}')
706
+ if image.shape[0] <= 0 or image.shape[1] <= 0:
707
+ raise ValueError(f'Invalid image shape={image.shape}')
708
+ if image.shape[2] != 3:
709
+ raise ValueError(f'Expected 3 channels, got shape={image.shape}')
710
+ if image.dtype != np.uint8:
711
+ image = image.astype(np.uint8)
712
+ input_tensor, ratio, pad, orig_size = self._preprocess(image)
713
+ expected = (1, 3, self.input_height, self.input_width)
714
+ if input_tensor.shape != expected:
715
+ raise ValueError(f'Bad input tensor shape={input_tensor.shape}, expected={expected}')
716
+ outputs = self.session.run(self.output_names, {self.input_name: input_tensor})
717
+ return self._postprocess(outputs[0], ratio, pad, orig_size)
718
+
719
+ def _predict_tta(self, image: np.ndarray) -> tuple[list[BoundingBox], tuple[np.ndarray, np.ndarray, np.ndarray]]:
720
+ boxes_orig, raw_orig = self._predict_single(image)
721
+ flipped = cv2.flip(image, 1)
722
+ boxes_flip, raw_flip = self._predict_single(flipped)
723
+ w = image.shape[1]
724
+ boxes_flip = [BoundingBox(x1=w - b.x2, y1=b.y1, x2=w - b.x1, y2=b.y2, cls_id=b.cls_id, conf=b.conf) for b in boxes_flip]
725
+ raw_boxes_o, raw_scores_o, raw_cls_o = raw_orig
726
+ raw_boxes_f, raw_scores_f, raw_cls_f = raw_flip
727
+ if len(raw_boxes_f) > 0:
728
+ mapped_f = raw_boxes_f.copy()
729
+ mapped_f[:, 0] = w - raw_boxes_f[:, 2]
730
+ mapped_f[:, 2] = w - raw_boxes_f[:, 0]
731
+ mapped_f[:, 1] = raw_boxes_f[:, 1]
732
+ mapped_f[:, 3] = raw_boxes_f[:, 3]
733
+ else:
734
+ mapped_f = raw_boxes_f
735
+ if len(raw_boxes_o) == 0 and len(mapped_f) == 0:
736
+ raw_all = self._empty_raw()
737
+ elif len(raw_boxes_o) == 0:
738
+ raw_all = (mapped_f, raw_scores_f, raw_cls_f)
739
+ elif len(mapped_f) == 0:
740
+ raw_all = (raw_boxes_o, raw_scores_o, raw_cls_o)
741
+ else:
742
+ raw_all = (np.concatenate([raw_boxes_o, mapped_f], axis=0), np.concatenate([raw_scores_o, raw_scores_f], axis=0), np.concatenate([raw_cls_o, raw_cls_f], axis=0))
743
+ all_boxes = boxes_orig + boxes_flip
744
+ if not all_boxes:
745
+ return ([], raw_all)
746
+ coords = np.array([[b.x1, b.y1, b.x2, b.y2] for b in all_boxes], dtype=np.float32)
747
+ scores = np.array([b.conf for b in all_boxes], dtype=np.float32)
748
+ cls_ids = np.array([b.cls_id for b in all_boxes], dtype=np.int32)
749
+ view_ids = np.array([0] * len(boxes_orig) + [1] * len(boxes_flip), dtype=np.int32)
750
+ hard_keep = self._per_class_hard_nms(coords, scores, cls_ids, self.iou_thres)
751
+ if len(hard_keep) == 0:
752
+ return ([], raw_all)
753
+ if len(hard_keep) > self.max_det:
754
+ top = np.argsort(-scores[hard_keep])[:self.max_det]
755
+ hard_keep = hard_keep[top]
756
+ boosted = self._max_score_per_cluster(coords[hard_keep], cls_ids[hard_keep], coords, scores, cls_ids, self.iou_thres)
757
+ kept_coords = coords[hard_keep]
758
+ kept_cls = cls_ids[hard_keep]
759
+ if self.use_tta_view_filter and self.tta_view_filter_max_conf > 0.0 and (len(kept_coords) > 0):
760
+ corrob = self._views_corroborated(kept_coords, kept_cls, coords, cls_ids, view_ids, self.tta_view_iou_thresh)
761
+ keep = ~((boosted <= self.tta_view_filter_max_conf) & ~corrob)
762
+ kept_coords = kept_coords[keep]
763
+ boosted = boosted[keep]
764
+ kept_cls = kept_cls[keep]
765
+ if len(kept_coords) > 1:
766
+ kept_coords, boosted, kept_cls = self._cross_class_dedup_op(kept_coords, boosted, kept_cls, self.cross_iou_thresh)
767
+ if len(kept_coords) > 1:
768
+ kept_coords, boosted, kept_cls = self._merge_same_class_boxes(kept_coords, boosted, kept_cls)
769
+ finals = [BoundingBox(x1=int(math.floor(kept_coords[j, 0])), y1=int(math.floor(kept_coords[j, 1])), x2=int(math.ceil(kept_coords[j, 2])), y2=int(math.ceil(kept_coords[j, 3])), cls_id=int(kept_cls[j]), conf=float(boosted[j])) for j in range(len(kept_coords))]
770
+ return (finals, raw_all)
771
+
772
+ def predict_batch(self, batch_images: list[ndarray], offset: int, n_keypoints: int) -> list[TVFrameResult]:
773
+ results: list[TVFrameResult] = []
774
+ for frame_number_in_batch, image in enumerate(batch_images):
775
+ try:
776
+ if self.use_tta:
777
+ boxes, raw = self._predict_tta(image)
778
+ else:
779
+ boxes, raw = self._predict_single(image)
780
+ if isinstance(image, np.ndarray) and image.ndim == 3:
781
+ boxes = self._filter_low_conf_by_color(image, boxes)
782
+ boxes = self._remove_edge_low_conf(boxes, (image.shape[1], image.shape[0]))
783
+ boxes = self._apply_extra_post(boxes, raw[0], raw[1], raw[2], (image.shape[1], image.shape[0]))
784
+ else:
785
+ boxes = self._apply_extra_post(boxes, raw[0], raw[1], raw[2], (0, 0))
786
+ except Exception as e:
787
+ print(f'⚠️ Inference failed for frame {offset + frame_number_in_batch}: {e}')
788
+ boxes = []
789
+ results.append(TVFrameResult(frame_id=offset + frame_number_in_batch, boxes=boxes, keypoints=[(0, 0) for _ in range(max(0, int(n_keypoints)))]))
790
+ return results
weights.onnx ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:c5b2e7c0bd4c1def3e715fdfa4c033a722f04605282d0ad5c91e840f7318fbc9
3
+ size 9734899