SuperBitDev commited on
Commit
1ef68ed
·
verified ·
1 Parent(s): d576b66

Upload folder using huggingface_hub

Browse files
Files changed (3) hide show
  1. chute_config.yml +21 -0
  2. miner.py +1106 -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>=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==2.8.0 torchvision==0.23.0 torchaudio==2.8.0 --index-url https://download.pytorch.org/whl/cu128
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,1106 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+
3
+ from pathlib import Path
4
+ import math
5
+
6
+ import cv2
7
+ import numpy as np
8
+ import onnxruntime as ort
9
+ from numpy import ndarray
10
+ from pydantic import BaseModel
11
+
12
+ class BoundingBox(BaseModel):
13
+ x1: int
14
+ y1: int
15
+ x2: int
16
+ y2: int
17
+ cls_id: int
18
+ conf: float
19
+
20
+
21
+ class TVFrameResult(BaseModel):
22
+ frame_id: int
23
+ boxes: list[BoundingBox]
24
+ keypoints: list[tuple[int, int]]
25
+
26
+
27
+ class Miner:
28
+ """ONNX Runtime miner for fire / smoke / fire_extinguisher detection.
29
+
30
+ Strategy (ported from offense miner):
31
+ - per-class confidence threshold with per-class rescue bonus
32
+ - per-class hard NMS, then cross-class dedup
33
+ - horizontal-flip TTA with full-set cluster score boost
34
+ Plus fire001 specifics: class remap, sanity-box filter, TTA toggle.
35
+ """
36
+
37
+ class_names = ["fire", "smoke", "fire extinguisher"]
38
+ # FALLBACK order the model emits classes in -- remapped to `class_names`
39
+ # index by `self.cls_remap` (built in __init__). The authoritative order
40
+ # is read from the ONNX `names` metadata that Ultralytics embeds at
41
+ # export time (ships inside weights.onnx), so a retrained model with a
42
+ # different class order is remapped correctly without code changes.
43
+ # Used only when that metadata is missing or unparsable.
44
+ _model_class_order = ["fire", "fire extinguisher", "smoke"]
45
+
46
+ iou_thres = 0.55
47
+ cross_iou_thresh = 0.8
48
+ max_det = 150
49
+
50
+ # Per-class confidence thresholds. Higher = fewer FP for that class.
51
+ # Indexed by class_names order: [fire, smoke, fire_extinguisher].
52
+ _conf_thres_array = np.array(
53
+ [0.18, 0.30, 0.2], dtype=np.float32
54
+ )
55
+ # Per-class rescue bonus. If a class has ZERO boxes passing the threshold
56
+ # in a frame, its top-1 candidate is admitted when its score is at least
57
+ # (threshold - bonus). Fire and smoke get a small bonus (variable
58
+ # appearance); fire extinguisher does not (distinctive object, leave FP
59
+ # control strict).
60
+ _bonus_array = np.array(
61
+ [0.03, 0.1, 0.05], dtype=np.float32
62
+ )
63
+
64
+ # Box sanity filter (fire001-specific FP reduction): drop tiny / degenerate
65
+ # / image-spanning / extreme aspect ratio boxes.
66
+ min_box_area = 14 * 14
67
+ min_side = 8
68
+ max_aspect_ratio = 8.0
69
+
70
+ # Same-class merge: two boxes whose intersection covers at least this
71
+ # fraction of the SMALLER box are treated as the same object and replaced
72
+ # by their union. Catches nested boxes (IoU below the NMS threshold) and
73
+ # fragmented detections. Per-class because the risk differs:
74
+ # smoke -- diffuse plumes fragment a lot, so a moderate threshold helps.
75
+ # fire -- separate flames must stay separate, so keep this HIGH (only a
76
+ # tight core nested inside a looser flame box merges). Set to a
77
+ # value > 1.0 to disable fire merging entirely.
78
+ # Fire merge is DISABLED by default (1.01): measured on the fire-29-val1024
79
+ # val split it cost fire AP (0.751 -> 0.742, composite 0.8888 -> 0.8874)
80
+ # because the nested core+flame boxes it collapses were scoring as separate
81
+ # true positives. Lower it to ~0.8 to enable, and re-measure with
82
+ # verify_filters.py / tune_miner.py after a retrain -- a model whose fire
83
+ # boxes fragment more (or live-SAM3 GT that draws fuller flames) could flip
84
+ # the result.
85
+ smoke_merge_overlap = 0.8
86
+ fire_merge_overlap = 1.01
87
+
88
+ # Fire containment suppression: when two FIRE boxes overlap on one object
89
+ # (intersection >= this fraction of the SMALLER box) keep the HIGHER-conf
90
+ # box and drop the other -- unchanged geometry, unlike the union merge
91
+ # above. This catches the nested core+flame duplicate that per-class NMS
92
+ # (IoU-based, iou_thres) leaves behind. Set > 1.0 to disable.
93
+ # DISABLED by default (1.01): measured on fire-29-val1024 it cost fire AP
94
+ # (0.751 -> 0.743, composite 0.8888 -> 0.8877). Cause: GT fire boxes almost
95
+ # never overlap (1 pair in 416), so each nested model pair has one TP + one
96
+ # FP, but the higher-CONF box isn't always the one matching GT at IoU 0.5 --
97
+ # so keeping it can drop the real match, and score-ordered AP already
98
+ # tolerates the duplicate. Lower to ~0.8 to enable; re-measure after a
99
+ # retrain or against live-SAM3 GT, which may differ.
100
+ fire_suppress_overlap = 0.88
101
+
102
+ # ── Low-confidence color-prior FP filters ───────────────────────────────
103
+ # Ported from the firedetect1007 miner's color checks, but applied ONLY to
104
+ # the borderline confidence band (just above each per-class threshold) and
105
+ # ONLY on color frames. A fire/extinguisher detection there is dropped when
106
+ # its pixels clearly do not match the expected appearance: warm/bright for
107
+ # fire, red for extinguisher. High-confidence detections are never touched.
108
+ #
109
+ # The reference miner ran these unconditionally -- a BUG on this validator,
110
+ # which feeds some frames as grayscale (a true red extinguisher is gray
111
+ # there, so a red test would wrongly delete it). We skip the filter when the
112
+ # ROI is near-grayscale, so it never fires on those frames.
113
+ #
114
+ # Tunable: set a max-conf gate to 0.0 to disable that filter. After a model
115
+ # retrain, re-validate these with tune_miner.py (the gates are relative to
116
+ # the per-class thresholds, so they move when those move).
117
+ fire_color_filter_max_conf = 0.45 # only fire boxes in (thresh, 0.45]
118
+ fire_ext_color_filter_max_conf = 0.40 # only ext boxes in (thresh, 0.40]
119
+ color_filter_min_saturation = 0.06 # skip filter if ROI is near-grayscale
120
+
121
+ # ── Corroboration FP filters (optional; OFF by default) ─────────────────
122
+ # Ported in spirit from firedetect1007. Both REMOVE borderline boxes that
123
+ # lack support -- a precision play for the validator's FP pillar. OFF by
124
+ # default because, unlike the color priors, they can also drop true
125
+ # positives; enable + sweep with verify_filters.py and keep only the
126
+ # settings that raise the measured composite. A max-conf gate of 0.0
127
+ # disables the corresponding filter.
128
+ # edge filter: drop boxes touching the frame border in a low-conf band
129
+ # (the validator scales/crops, so border-hugging boxes are often the
130
+ # truncated remains of an object whose body is off-frame).
131
+ # tta view filter: drop low-conf boxes that appear in only ONE of the two
132
+ # horizontal-flip TTA views (a real object is usually seen in both).
133
+ use_edge_filter = False
134
+ edge_filter_max_conf = 0.0 # drop edge-touching boxes with conf <= this
135
+ edge_tol = 2.0 # px from the border counted as "on edge"
136
+ use_tta_view_filter = False
137
+ tta_view_filter_max_conf = 0.0 # drop single-view boxes with conf <= this
138
+ tta_view_iou_thresh = 0.5 # IoU for "same object seen in both views"
139
+
140
+ def __init__(self, path_hf_repo: Path) -> None:
141
+ model_path = path_hf_repo / "weights.onnx"
142
+ print("ORT version:", ort.__version__)
143
+
144
+ try:
145
+ ort.preload_dlls()
146
+ print("✅ onnxruntime.preload_dlls() success")
147
+ except Exception as e:
148
+ print(f"⚠️ preload_dlls failed: {e}")
149
+
150
+ print("ORT available providers BEFORE session:", ort.get_available_providers())
151
+
152
+ sess_options = ort.SessionOptions()
153
+ sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
154
+ sess_options.intra_op_num_threads = 2
155
+ sess_options.inter_op_num_threads = 1
156
+ sess_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
157
+
158
+ try:
159
+ self.session = ort.InferenceSession(
160
+ str(model_path),
161
+ sess_options=sess_options,
162
+ providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
163
+ )
164
+ print("✅ Created ORT session with preferred CUDA provider list")
165
+ except Exception as e:
166
+ print(f"⚠️ CUDA session creation failed, falling back to CPU: {e}")
167
+ self.session = ort.InferenceSession(
168
+ str(model_path),
169
+ sess_options=sess_options,
170
+ providers=["CPUExecutionProvider"],
171
+ )
172
+
173
+ print("ORT session providers:", self.session.get_providers())
174
+
175
+ # Build cls_remap: for each model-emit index i,
176
+ # cls_remap[i] = self.class_names.index(model_class_order[i])
177
+ # i.e. converts a model-side class id into the canonical class id
178
+ # that downstream code (BoundingBox.cls_id, validator) expects.
179
+ # The model-side order comes from the ONNX metadata when available,
180
+ # else falls back to the static _model_class_order.
181
+ model_class_order = self._read_model_class_order()
182
+ if model_class_order is None:
183
+ model_class_order = list(self._model_class_order)
184
+ print(f"cls order: no usable ONNX metadata, FALLBACK {model_class_order}")
185
+ else:
186
+ print(f"cls order: from ONNX metadata {model_class_order}")
187
+ self.cls_remap = np.array(
188
+ [self.class_names.index(n) for n in model_class_order],
189
+ dtype=np.int32,
190
+ )
191
+
192
+ for inp in self.session.get_inputs():
193
+ print("INPUT:", inp.name, inp.shape, inp.type)
194
+ for out in self.session.get_outputs():
195
+ print("OUTPUT:", out.name, out.shape, out.type)
196
+
197
+ self.input_name = self.session.get_inputs()[0].name
198
+ self.output_names = [output.name for output in self.session.get_outputs()]
199
+ self.input_shape = self.session.get_inputs()[0].shape
200
+
201
+ self.input_height = self._safe_dim(self.input_shape[2], default=1280)
202
+ self.input_width = self._safe_dim(self.input_shape[3], default=1280)
203
+
204
+ self.use_tta = False
205
+
206
+ print(f"✅ ONNX model loaded from: {model_path}")
207
+ print(f"✅ ONNX providers: {self.session.get_providers()}")
208
+ print(f"✅ ONNX input: name={self.input_name}, shape={self.input_shape}")
209
+ print("per-class conf: " + ", ".join(
210
+ f"{n}={t:.3f}" for n, t in zip(
211
+ self.class_names, self._conf_thres_array.tolist()
212
+ )
213
+ ))
214
+
215
+ self._warmup()
216
+
217
+ def _warmup(self, iters: int = 3) -> None:
218
+ try:
219
+ dummy = np.zeros((720, 1280, 3), dtype=np.uint8)
220
+ for _ in range(max(1, iters)):
221
+ self.predict_batch(batch_images=[dummy], offset=0, n_keypoints=0)
222
+ print(f"✅ warmup: {iters} dummy predict_batch call(s) done")
223
+ except Exception as e:
224
+ print(f"⚠️ warmup skipped: {e}")
225
+
226
+ def __repr__(self) -> str:
227
+ return (
228
+ f"ONNXRuntime(session={type(self.session).__name__}, "
229
+ f"providers={self.session.get_providers()})"
230
+ )
231
+
232
+ @staticmethod
233
+ def _safe_dim(value, default: int) -> int:
234
+ return value if isinstance(value, int) and value > 0 else default
235
+
236
+ def _read_model_class_order(self) -> list[str] | None:
237
+ """Read the model's class order from Ultralytics ONNX metadata.
238
+
239
+ Returns the class names ordered by model-emit index, or None when
240
+ metadata is missing/unparsable or doesn't match `class_names` as a
241
+ set (in which case the static _model_class_order fallback is used).
242
+ """
243
+ try:
244
+ import ast
245
+
246
+ meta = self.session.get_modelmeta().custom_metadata_map
247
+ names = ast.literal_eval(meta["names"]) # e.g. {0: 'fire', ...}
248
+ if isinstance(names, dict):
249
+ order = [str(names[i]) for i in sorted(names)]
250
+ else:
251
+ order = [str(n) for n in names]
252
+ except Exception as e:
253
+ print(f"cls order: could not read ONNX names metadata ({e})")
254
+ return None
255
+ if sorted(order) != sorted(self.class_names):
256
+ print(
257
+ f"cls order: ONNX names {order} do not match expected classes "
258
+ f"{self.class_names}; ignoring metadata"
259
+ )
260
+ return None
261
+ return order
262
+
263
+ def _letterbox(
264
+ self,
265
+ image: ndarray,
266
+ new_shape: tuple[int, int],
267
+ color=(114, 114, 114),
268
+ ) -> tuple[ndarray, float, tuple[float, float]]:
269
+ h, w = image.shape[:2]
270
+ new_w, new_h = new_shape
271
+
272
+ ratio = min(new_w / w, new_h / h)
273
+ resized_w = int(round(w * ratio))
274
+ resized_h = int(round(h * ratio))
275
+
276
+ if (resized_w, resized_h) != (w, h):
277
+ interp = cv2.INTER_CUBIC if ratio > 1.0 else cv2.INTER_LINEAR
278
+ image = cv2.resize(image, (resized_w, resized_h), interpolation=interp)
279
+
280
+ dw = (new_w - resized_w) / 2.0
281
+ dh = (new_h - resized_h) / 2.0
282
+
283
+ left = int(round(dw - 0.1))
284
+ right = int(round(dw + 0.1))
285
+ top = int(round(dh - 0.1))
286
+ bottom = int(round(dh + 0.1))
287
+
288
+ padded = cv2.copyMakeBorder(
289
+ image, top, bottom, left, right,
290
+ borderType=cv2.BORDER_CONSTANT, value=color,
291
+ )
292
+ return padded, ratio, (dw, dh)
293
+
294
+ def _preprocess(
295
+ self, image: ndarray
296
+ ) -> tuple[np.ndarray, float, tuple[float, float], tuple[int, int]]:
297
+ orig_h, orig_w = image.shape[:2]
298
+ img, ratio, pad = self._letterbox(
299
+ image, (self.input_width, self.input_height)
300
+ )
301
+ # Fused scale(1/255) + BGR->RGB swap + HWC->NCHW + contiguous float32 in
302
+ # one optimized OpenCV call. Bit-identical (max abs diff 6e-8) to the
303
+ # prior cvtColor + astype/255 + transpose + ascontiguousarray chain, but
304
+ # ~half the preprocess time (preprocess is ~12% of predict_batch).
305
+ blob = cv2.dnn.blobFromImage(img, scalefactor=1.0 / 255.0, swapRB=True)
306
+ return blob, ratio, pad, (orig_w, orig_h)
307
+
308
+ @staticmethod
309
+ def _clip_boxes(boxes: np.ndarray, image_size: tuple[int, int]) -> np.ndarray:
310
+ w, h = image_size
311
+ boxes[:, 0] = np.clip(boxes[:, 0], 0, w - 1)
312
+ boxes[:, 1] = np.clip(boxes[:, 1], 0, h - 1)
313
+ boxes[:, 2] = np.clip(boxes[:, 2], 0, w - 1)
314
+ boxes[:, 3] = np.clip(boxes[:, 3], 0, h - 1)
315
+ return boxes
316
+
317
+ @staticmethod
318
+ def _xywh_to_xyxy(boxes: np.ndarray) -> np.ndarray:
319
+ out = np.empty_like(boxes)
320
+ out[:, 0] = boxes[:, 0] - boxes[:, 2] / 2.0
321
+ out[:, 1] = boxes[:, 1] - boxes[:, 3] / 2.0
322
+ out[:, 2] = boxes[:, 0] + boxes[:, 2] / 2.0
323
+ out[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.0
324
+ return out
325
+
326
+ @staticmethod
327
+ def _hard_nms(
328
+ boxes: np.ndarray, scores: np.ndarray, iou_thresh: float
329
+ ) -> np.ndarray:
330
+ n = len(boxes)
331
+ if n == 0:
332
+ return np.array([], dtype=np.intp)
333
+ order = np.argsort(-scores)
334
+ keep: list[int] = []
335
+ while len(order) > 0:
336
+ i = int(order[0])
337
+ keep.append(i)
338
+ if len(order) == 1:
339
+ break
340
+ rest = order[1:]
341
+ xx1 = np.maximum(boxes[i, 0], boxes[rest, 0])
342
+ yy1 = np.maximum(boxes[i, 1], boxes[rest, 1])
343
+ xx2 = np.minimum(boxes[i, 2], boxes[rest, 2])
344
+ yy2 = np.minimum(boxes[i, 3], boxes[rest, 3])
345
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
346
+ a_i = (max(0.0, boxes[i, 2] - boxes[i, 0]) *
347
+ max(0.0, boxes[i, 3] - boxes[i, 1]))
348
+ a_r = (np.maximum(0.0, boxes[rest, 2] - boxes[rest, 0]) *
349
+ np.maximum(0.0, boxes[rest, 3] - boxes[rest, 1]))
350
+ iou = inter / (a_i + a_r - inter + 1e-7)
351
+ order = rest[iou <= iou_thresh]
352
+ return np.array(keep, dtype=np.intp)
353
+
354
+ def _per_class_hard_nms(
355
+ self,
356
+ boxes: np.ndarray,
357
+ scores: np.ndarray,
358
+ cls_ids: np.ndarray,
359
+ iou_thresh: float,
360
+ ) -> np.ndarray:
361
+ if len(boxes) == 0:
362
+ return np.array([], dtype=np.intp)
363
+ all_keep: list[int] = []
364
+ for c in np.unique(cls_ids):
365
+ mask = cls_ids == c
366
+ indices = np.where(mask)[0]
367
+ keep = self._hard_nms(boxes[mask], scores[mask], iou_thresh)
368
+ all_keep.extend(indices[keep].tolist())
369
+ all_keep.sort()
370
+ return np.array(all_keep, dtype=np.intp)
371
+
372
+ def _cross_class_dedup_op(
373
+ self,
374
+ boxes: np.ndarray,
375
+ scores: np.ndarray,
376
+ cls_ids: np.ndarray,
377
+ iou_thresh: float,
378
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
379
+ """Remove near-duplicate boxes across classes.
380
+
381
+ Order candidates by (score - per_class_threshold) margin, then by area;
382
+ keep the highest, suppress every other box with IoU > iou_thresh.
383
+ This suppresses the case where the same physical object is detected
384
+ as multiple classes (e.g. fire vs smoke on the same flames).
385
+ """
386
+ n = len(boxes)
387
+ if n <= 1:
388
+ return boxes, scores, cls_ids
389
+ boxes = np.asarray(boxes, dtype=np.float32)
390
+ scores = np.asarray(scores, dtype=np.float32)
391
+ cls_ids = np.asarray(cls_ids, dtype=np.int32)
392
+ areas = (np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) *
393
+ np.maximum(0.0, boxes[:, 3] - boxes[:, 1]))
394
+ margins = scores - self._conf_thres_array[cls_ids]
395
+ order = np.lexsort((-areas, -margins))
396
+ suppressed = np.zeros(n, dtype=bool)
397
+ keep: list[int] = []
398
+ for i in order:
399
+ if suppressed[i]:
400
+ continue
401
+ keep.append(int(i))
402
+ bi = boxes[i]
403
+ xx1 = np.maximum(bi[0], boxes[:, 0])
404
+ yy1 = np.maximum(bi[1], boxes[:, 1])
405
+ xx2 = np.minimum(bi[2], boxes[:, 2])
406
+ yy2 = np.minimum(bi[3], boxes[:, 3])
407
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
408
+ a_i = max(1e-7, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
409
+ iou = inter / (a_i + areas - inter + 1e-7)
410
+ dup = iou > iou_thresh
411
+ dup[i] = False
412
+ suppressed |= dup
413
+ keep_idx = np.array(keep, dtype=np.intp)
414
+ return boxes[keep_idx], scores[keep_idx], cls_ids[keep_idx]
415
+
416
+ def _merge_class_boxes(
417
+ self,
418
+ boxes: np.ndarray,
419
+ scores: np.ndarray,
420
+ cls_ids: np.ndarray,
421
+ target_cls: int,
422
+ overlap: float,
423
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
424
+ """Merge overlapping detections of ONE class into single boxes.
425
+
426
+ Two same-class boxes whose intersection covers >= `overlap` of the
427
+ SMALLER box are treated as one object and replaced by their union with
428
+ the max confidence of the pair. Repeats until no pair merges, so chains
429
+ of fragments collapse. `overlap` is intersection-over-minimum-area, so
430
+ only nested / heavily-overlapping boxes merge -- two spatially separate
431
+ objects (low mutual overlap) are never fused. `overlap > 1.0` disables.
432
+ """
433
+ if overlap > 1.0:
434
+ return boxes, scores, cls_ids
435
+ idx = np.where(cls_ids == target_cls)[0]
436
+ if len(idx) <= 1:
437
+ return boxes, scores, cls_ids
438
+
439
+ sb = boxes[idx].astype(np.float32).tolist()
440
+ ss = scores[idx].astype(np.float32).tolist()
441
+ merged_any = True
442
+ while merged_any and len(sb) > 1:
443
+ merged_any = False
444
+ for i in range(len(sb)):
445
+ for j in range(i + 1, len(sb)):
446
+ a, b = sb[i], sb[j]
447
+ ix1 = max(a[0], b[0])
448
+ iy1 = max(a[1], b[1])
449
+ ix2 = min(a[2], b[2])
450
+ iy2 = min(a[3], b[3])
451
+ inter = max(0.0, ix2 - ix1) * max(0.0, iy2 - iy1)
452
+ area_a = max(0.0, a[2] - a[0]) * max(0.0, a[3] - a[1])
453
+ area_b = max(0.0, b[2] - b[0]) * max(0.0, b[3] - b[1])
454
+ smaller = min(area_a, area_b)
455
+ if inter / (smaller + 1e-7) >= overlap:
456
+ sb[i] = [
457
+ min(a[0], b[0]), min(a[1], b[1]),
458
+ max(a[2], b[2]), max(a[3], b[3]),
459
+ ]
460
+ ss[i] = max(ss[i], ss[j])
461
+ del sb[j]
462
+ del ss[j]
463
+ merged_any = True
464
+ break
465
+ if merged_any:
466
+ break
467
+
468
+ other = cls_ids != target_cls
469
+ new_boxes = np.concatenate(
470
+ [boxes[other].astype(np.float32),
471
+ np.array(sb, dtype=np.float32).reshape(-1, 4)]
472
+ )
473
+ new_scores = np.concatenate(
474
+ [scores[other].astype(np.float32),
475
+ np.array(ss, dtype=np.float32)]
476
+ )
477
+ new_cls = np.concatenate(
478
+ [cls_ids[other].astype(np.int32),
479
+ np.full(len(sb), target_cls, dtype=np.int32)]
480
+ )
481
+ return new_boxes, new_scores, new_cls
482
+
483
+ def _suppress_contained_lower_conf(
484
+ self,
485
+ boxes: np.ndarray,
486
+ scores: np.ndarray,
487
+ cls_ids: np.ndarray,
488
+ target_cls: int,
489
+ overlap: float,
490
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
491
+ """For one class, when two boxes overlap (intersection >= `overlap` of
492
+ the smaller box) keep the higher-confidence box and drop the other.
493
+ Geometry is never changed -- only the redundant lower-conf box is
494
+ removed. `overlap > 1.0` disables."""
495
+ if overlap > 1.0:
496
+ return boxes, scores, cls_ids
497
+ idx = np.where(cls_ids == target_cls)[0]
498
+ if len(idx) <= 1:
499
+ return boxes, scores, cls_ids
500
+
501
+ order = idx[np.argsort(-scores[idx])] # highest confidence first
502
+ remove: set[int] = set()
503
+ for a in range(len(order)):
504
+ i = int(order[a])
505
+ if i in remove:
506
+ continue
507
+ bi = boxes[i]
508
+ area_i = max(1e-7, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
509
+ for b in range(a + 1, len(order)):
510
+ j = int(order[b])
511
+ if j in remove:
512
+ continue
513
+ bj = boxes[j]
514
+ ix1 = max(bi[0], bj[0]); iy1 = max(bi[1], bj[1])
515
+ ix2 = min(bi[2], bj[2]); iy2 = min(bi[3], bj[3])
516
+ inter = max(0.0, ix2 - ix1) * max(0.0, iy2 - iy1)
517
+ if inter <= 0.0:
518
+ continue
519
+ area_j = max(1e-7, float((bj[2] - bj[0]) * (bj[3] - bj[1])))
520
+ if inter / (min(area_i, area_j) + 1e-7) >= overlap:
521
+ remove.add(j) # j is the lower-confidence box (order desc)
522
+ if not remove:
523
+ return boxes, scores, cls_ids
524
+ keep = np.array(
525
+ [k not in remove for k in range(len(boxes))], dtype=bool
526
+ )
527
+ return boxes[keep], scores[keep], cls_ids[keep]
528
+
529
+ def _merge_same_class_boxes(
530
+ self,
531
+ boxes: np.ndarray,
532
+ scores: np.ndarray,
533
+ cls_ids: np.ndarray,
534
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
535
+ """Resolve nested / fragmented same-object detections, per class.
536
+
537
+ Smoke: diffuse plumes fragment into nested boxes NMS can't collapse, so
538
+ they are UNION-merged (smoke_merge_overlap).
539
+ Fire: a tight hot-core box and a looser flame box are the same flame;
540
+ keep the HIGHER-confidence one and drop the other (fire_suppress_overlap),
541
+ which leaves geometry intact. The union-merge variant (fire_merge_overlap)
542
+ is also available but measured worse, so it is disabled by default.
543
+ """
544
+ boxes, scores, cls_ids = self._merge_class_boxes(
545
+ boxes, scores, cls_ids,
546
+ self.class_names.index("smoke"), self.smoke_merge_overlap,
547
+ )
548
+ boxes, scores, cls_ids = self._merge_class_boxes(
549
+ boxes, scores, cls_ids,
550
+ self.class_names.index("fire"), self.fire_merge_overlap,
551
+ )
552
+ boxes, scores, cls_ids = self._suppress_contained_lower_conf(
553
+ boxes, scores, cls_ids,
554
+ self.class_names.index("fire"), self.fire_suppress_overlap,
555
+ )
556
+ return boxes, scores, cls_ids
557
+
558
+ # Back-compat alias (older callers / tune_miner referenced this name).
559
+ def _merge_smoke_boxes(
560
+ self,
561
+ boxes: np.ndarray,
562
+ scores: np.ndarray,
563
+ cls_ids: np.ndarray,
564
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
565
+ return self._merge_same_class_boxes(boxes, scores, cls_ids)
566
+
567
+ @staticmethod
568
+ def _max_score_per_cluster(
569
+ post_boxes: np.ndarray,
570
+ post_cls: np.ndarray,
571
+ full_boxes: np.ndarray,
572
+ full_scores: np.ndarray,
573
+ full_cls: np.ndarray,
574
+ iou_thresh: float,
575
+ ) -> np.ndarray:
576
+ """For each kept (post-NMS) box, return the max score over the FULL
577
+ candidate set among same-class boxes with IoU >= iou_thresh.
578
+
579
+ Used after horizontal-flip TTA: a high-confidence flipped detection
580
+ can raise the score of the corresponding original detection.
581
+ """
582
+ n = len(post_boxes)
583
+ if n == 0:
584
+ return np.empty(0, dtype=np.float32)
585
+ full_areas = (np.maximum(0.0, full_boxes[:, 2] - full_boxes[:, 0]) *
586
+ np.maximum(0.0, full_boxes[:, 3] - full_boxes[:, 1]))
587
+ out = np.empty(n, dtype=np.float32)
588
+ for i in range(n):
589
+ bi = post_boxes[i]
590
+ xx1 = np.maximum(bi[0], full_boxes[:, 0])
591
+ yy1 = np.maximum(bi[1], full_boxes[:, 1])
592
+ xx2 = np.minimum(bi[2], full_boxes[:, 2])
593
+ yy2 = np.minimum(bi[3], full_boxes[:, 3])
594
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
595
+ a_i = max(0.0, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
596
+ iou = inter / (a_i + full_areas - inter + 1e-7)
597
+ cluster = (iou >= iou_thresh) & (full_cls == post_cls[i])
598
+ out[i] = float(np.max(full_scores[cluster])) if np.any(cluster) else 0.0
599
+ return out
600
+
601
+ def _conf_filter_mask(
602
+ self, scores: np.ndarray, cls_ids: np.ndarray
603
+ ) -> np.ndarray:
604
+ """Boolean keep-mask: score >= per-class threshold, with a per-class
605
+ rescue -- if a class has zero boxes passing, admit its top-1 candidate
606
+ when its score >= (per-class threshold - per-class bonus)."""
607
+ if len(scores) == 0:
608
+ return np.zeros(0, dtype=bool)
609
+ thr = self._conf_thres_array[cls_ids]
610
+ keep = scores >= thr
611
+ for c in np.unique(cls_ids):
612
+ b = float(self._bonus_array[c])
613
+ if b <= 0.0:
614
+ continue
615
+ cm = cls_ids == c
616
+ if keep[cm].any():
617
+ continue
618
+ idx = np.where(cm)[0]
619
+ top = int(idx[int(np.argmax(scores[idx]))])
620
+ if scores[top] >= self._conf_thres_array[c] - b:
621
+ keep[top] = True
622
+ return keep
623
+
624
+ def _filter_sane_boxes(
625
+ self,
626
+ boxes: np.ndarray,
627
+ scores: np.ndarray,
628
+ cls_ids: np.ndarray,
629
+ orig_size: tuple[int, int],
630
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
631
+ """Drop tiny / degenerate / image-spanning / extreme-AR boxes (FP)."""
632
+ if len(boxes) == 0:
633
+ return boxes, scores, cls_ids
634
+ orig_w, orig_h = orig_size
635
+ image_area = float(orig_w * orig_h)
636
+ keep = []
637
+ for i, box in enumerate(boxes):
638
+ x1, y1, x2, y2 = box.tolist()
639
+ bw = x2 - x1
640
+ bh = y2 - y1
641
+ if bw <= 0 or bh <= 0:
642
+ continue
643
+ if bw < self.min_side or bh < self.min_side:
644
+ continue
645
+ area = bw * bh
646
+ if area < self.min_box_area:
647
+ continue
648
+ if area > 0.95 * image_area:
649
+ continue
650
+ ar = max(bw / max(bh, 1e-6), bh / max(bw, 1e-6))
651
+ if ar > self.max_aspect_ratio:
652
+ continue
653
+ keep.append(i)
654
+ if not keep:
655
+ return (
656
+ np.empty((0, 4), dtype=np.float32),
657
+ np.empty((0,), dtype=np.float32),
658
+ np.empty((0,), dtype=np.int32),
659
+ )
660
+ k = np.array(keep, dtype=np.intp)
661
+ return boxes[k], scores[k], cls_ids[k]
662
+
663
+ def _per_view_pipeline(
664
+ self,
665
+ boxes: np.ndarray,
666
+ scores: np.ndarray,
667
+ cls_ids: np.ndarray,
668
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
669
+ """Per-view post-processing pipeline: per-class NMS -> cap -> cross-class dedup -> smoke merge."""
670
+ if len(boxes) > 1:
671
+ keep = self._per_class_hard_nms(boxes, scores, cls_ids, self.iou_thres)
672
+ boxes, scores, cls_ids = boxes[keep], scores[keep], cls_ids[keep]
673
+ if len(scores) > self.max_det:
674
+ top = np.argsort(-scores)[: self.max_det]
675
+ boxes, scores, cls_ids = boxes[top], scores[top], cls_ids[top]
676
+ if len(boxes) > 1:
677
+ boxes, scores, cls_ids = self._cross_class_dedup_op(
678
+ boxes, scores, cls_ids, self.cross_iou_thresh
679
+ )
680
+ if len(boxes) > 1:
681
+ boxes, scores, cls_ids = self._merge_same_class_boxes(boxes, scores, cls_ids)
682
+ return boxes, scores, cls_ids
683
+
684
+ @staticmethod
685
+ def _roi_for_box(image: np.ndarray, box: BoundingBox) -> np.ndarray | None:
686
+ """Clip a BoundingBox to the image and return its BGR pixel ROI."""
687
+ h, w = image.shape[:2]
688
+ x1 = max(0, int(math.floor(box.x1)))
689
+ y1 = max(0, int(math.floor(box.y1)))
690
+ x2 = min(w, int(math.ceil(box.x2)))
691
+ y2 = min(h, int(math.ceil(box.y2)))
692
+ if x2 <= x1 or y2 <= y1:
693
+ return None
694
+ roi = image[y1:y2, x1:x2]
695
+ return roi if roi.size else None
696
+
697
+ def _roi_is_near_grayscale(self, roi: np.ndarray) -> bool:
698
+ """True if the ROI carries almost no color (validator grayscale frame).
699
+ On such ROIs the color priors are skipped so they can't delete valid
700
+ red/warm objects that have been stripped of color."""
701
+ mx = roi.max(axis=2).astype(np.float32)
702
+ mn = roi.min(axis=2).astype(np.float32)
703
+ sat = (mx - mn) / (mx + 1e-6)
704
+ return float(sat.mean()) < self.color_filter_min_saturation
705
+
706
+ @staticmethod
707
+ def _passes_fire_color(roi: np.ndarray) -> bool:
708
+ """Fire is warm and/or has a bright hotspot. ROI is BGR."""
709
+ blue = roi[:, :, 0].astype(np.float32)
710
+ green = roi[:, :, 1].astype(np.float32)
711
+ red = roi[:, :, 2].astype(np.float32)
712
+ mean_r = float(np.mean(red))
713
+ max_rgb = float(max(np.max(red), np.max(green), np.max(blue)))
714
+ bright_frac = float(np.mean(np.max(roi, axis=2) >= 150))
715
+ # A bright hotspot is fire-like even with little hue (also covers the
716
+ # near-white core of an intense flame).
717
+ if max_rgb >= 200.0 and bright_frac >= 0.01:
718
+ return True
719
+ warm = (red > green + 10.0) & (red > blue + 10.0)
720
+ warm_frac = float(np.mean(warm))
721
+ r_minus_g = mean_r - float(np.mean(green))
722
+ if warm_frac >= 0.05 and (
723
+ max_rgb >= 120.0 or mean_r >= 120.0 or warm_frac >= 0.15
724
+ ):
725
+ return True
726
+ if bright_frac >= 0.12 and r_minus_g >= 2.0:
727
+ return True
728
+ return False
729
+
730
+ @staticmethod
731
+ def _passes_fire_ext_red_color(roi: np.ndarray) -> bool:
732
+ """Fire extinguishers are red. ROI is BGR. Lenient: only clearly
733
+ cool/green/blue or very dark regions fail."""
734
+ blue = roi[:, :, 0].astype(np.float32)
735
+ green = roi[:, :, 1].astype(np.float32)
736
+ red = roi[:, :, 2].astype(np.float32)
737
+ red_dom = float(np.mean((red > green + 10.0) & (red > blue + 10.0)))
738
+ if red_dom >= 0.03:
739
+ return True
740
+ if (float(np.mean(red)) - float(np.mean(green))) >= 0.0 and \
741
+ float(np.mean(red)) >= 50.0:
742
+ return True
743
+ return False
744
+
745
+ def _remove_edge_low_conf(
746
+ self, results: list[BoundingBox], orig_size: tuple[int, int]
747
+ ) -> list[BoundingBox]:
748
+ """Drop border-hugging boxes in the low-confidence band."""
749
+ if (
750
+ not self.use_edge_filter
751
+ or self.edge_filter_max_conf <= 0.0
752
+ or not results
753
+ ):
754
+ return results
755
+ w, h = orig_size
756
+ tol = self.edge_tol
757
+ out: list[BoundingBox] = []
758
+ for b in results:
759
+ on_edge = (
760
+ b.x1 <= tol
761
+ or b.y1 <= tol
762
+ or b.x2 >= w - 1 - tol
763
+ or b.y2 >= h - 1 - tol
764
+ )
765
+ if on_edge and b.conf <= self.edge_filter_max_conf:
766
+ continue
767
+ out.append(b)
768
+ return out
769
+
770
+ def _views_corroborated(
771
+ self,
772
+ post_boxes: np.ndarray,
773
+ post_cls: np.ndarray,
774
+ full_boxes: np.ndarray,
775
+ full_cls: np.ndarray,
776
+ full_views: np.ndarray,
777
+ iou_thresh: float,
778
+ ) -> np.ndarray:
779
+ """For each post-NMS box, True if same-class detections from >= 2
780
+ distinct TTA views overlap it (IoU >= iou_thresh) in the full union."""
781
+ n = len(post_boxes)
782
+ if n == 0:
783
+ return np.zeros(0, dtype=bool)
784
+ full_areas = (np.maximum(0.0, full_boxes[:, 2] - full_boxes[:, 0]) *
785
+ np.maximum(0.0, full_boxes[:, 3] - full_boxes[:, 1]))
786
+ out = np.zeros(n, dtype=bool)
787
+ for i in range(n):
788
+ bi = post_boxes[i]
789
+ xx1 = np.maximum(bi[0], full_boxes[:, 0])
790
+ yy1 = np.maximum(bi[1], full_boxes[:, 1])
791
+ xx2 = np.minimum(bi[2], full_boxes[:, 2])
792
+ yy2 = np.minimum(bi[3], full_boxes[:, 3])
793
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
794
+ a_i = max(0.0, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
795
+ iou = inter / (a_i + full_areas - inter + 1e-7)
796
+ mask = (iou >= iou_thresh) & (full_cls == post_cls[i])
797
+ if np.any(mask):
798
+ out[i] = len(np.unique(full_views[mask])) >= 2
799
+ return out
800
+
801
+ def _filter_low_conf_by_color(
802
+ self, image: np.ndarray, results: list[BoundingBox]
803
+ ) -> list[BoundingBox]:
804
+ """Drop borderline fire / extinguisher detections whose pixels clearly
805
+ contradict the class's expected color. No-op on near-grayscale ROIs and
806
+ on detections above the per-class color-filter conf gate."""
807
+ if not results:
808
+ return results
809
+ cls_fire = self.class_names.index("fire")
810
+ cls_ext = self.class_names.index("fire extinguisher")
811
+ out: list[BoundingBox] = []
812
+ for box in results:
813
+ check_fire = (
814
+ box.cls_id == cls_fire
815
+ and box.conf <= self.fire_color_filter_max_conf
816
+ )
817
+ check_ext = (
818
+ box.cls_id == cls_ext
819
+ and box.conf <= self.fire_ext_color_filter_max_conf
820
+ )
821
+ if not check_fire and not check_ext:
822
+ out.append(box)
823
+ continue
824
+ roi = self._roi_for_box(image, box)
825
+ if roi is None or self._roi_is_near_grayscale(roi):
826
+ out.append(box)
827
+ continue
828
+ if check_fire and not self._passes_fire_color(roi):
829
+ continue
830
+ if check_ext and not self._passes_fire_ext_red_color(roi):
831
+ continue
832
+ out.append(box)
833
+ return out
834
+
835
+ @staticmethod
836
+ def _build_results(
837
+ boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray
838
+ ) -> list[BoundingBox]:
839
+ results: list[BoundingBox] = []
840
+ for box, conf, cls_id in zip(boxes, scores, cls_ids):
841
+ x1, y1, x2, y2 = box.tolist()
842
+ if x2 <= x1 or y2 <= y1:
843
+ continue
844
+ results.append(
845
+ BoundingBox(
846
+ x1=int(math.floor(x1)),
847
+ y1=int(math.floor(y1)),
848
+ x2=int(math.ceil(x2)),
849
+ y2=int(math.ceil(y2)),
850
+ cls_id=int(cls_id),
851
+ conf=float(conf),
852
+ )
853
+ )
854
+ return results
855
+
856
+ def _decode_final_dets(
857
+ self,
858
+ preds: np.ndarray,
859
+ ratio: float,
860
+ pad: tuple[float, float],
861
+ orig_size: tuple[int, int],
862
+ ) -> list[BoundingBox]:
863
+ """Final-detection output path: rows shaped [x1, y1, x2, y2, conf, cls_id]."""
864
+ if preds.ndim == 3 and preds.shape[0] == 1:
865
+ preds = preds[0]
866
+ if preds.ndim != 2 or preds.shape[1] < 6:
867
+ raise ValueError(f"Unexpected ONNX final-det output shape: {preds.shape}")
868
+
869
+ boxes = preds[:, :4].astype(np.float32)
870
+ scores = preds[:, 4].astype(np.float32)
871
+ cls_ids = preds[:, 5].astype(np.int32)
872
+ cls_ids = self.cls_remap[cls_ids]
873
+
874
+ keep = self._conf_filter_mask(scores, cls_ids)
875
+ boxes = boxes[keep]
876
+ scores = scores[keep]
877
+ cls_ids = cls_ids[keep]
878
+ if len(boxes) == 0:
879
+ return []
880
+
881
+ pad_w, pad_h = pad
882
+ boxes[:, [0, 2]] -= pad_w
883
+ boxes[:, [1, 3]] -= pad_h
884
+ boxes /= ratio
885
+ boxes = self._clip_boxes(boxes, orig_size)
886
+
887
+ boxes, scores, cls_ids = self._filter_sane_boxes(
888
+ boxes, scores, cls_ids, orig_size
889
+ )
890
+ if len(boxes) == 0:
891
+ return []
892
+
893
+ boxes, scores, cls_ids = self._per_view_pipeline(boxes, scores, cls_ids)
894
+ return self._build_results(boxes, scores, cls_ids)
895
+
896
+ def _decode_raw_yolo(
897
+ self,
898
+ preds: np.ndarray,
899
+ ratio: float,
900
+ pad: tuple[float, float],
901
+ orig_size: tuple[int, int],
902
+ ) -> list[BoundingBox]:
903
+ """Fallback raw-YOLO output path: per-anchor class logits."""
904
+ if preds.ndim != 3 or preds.shape[0] != 1:
905
+ raise ValueError(f"Unexpected raw ONNX output shape: {preds.shape}")
906
+ preds = preds[0]
907
+ if preds.shape[0] <= 16 and preds.shape[1] > preds.shape[0]:
908
+ preds = preds.T
909
+ if preds.ndim != 2 or preds.shape[1] < 5:
910
+ raise ValueError(f"Unexpected raw output shape: {preds.shape}")
911
+
912
+ boxes_xywh = preds[:, :4].astype(np.float32)
913
+ cls_part = preds[:, 4:].astype(np.float32)
914
+ if cls_part.shape[1] == 1:
915
+ scores = cls_part[:, 0]
916
+ cls_ids = np.zeros(len(scores), dtype=np.int32)
917
+ else:
918
+ cls_ids = np.argmax(cls_part, axis=1).astype(np.int32)
919
+ scores = cls_part[np.arange(len(cls_part)), cls_ids]
920
+ cls_ids = self.cls_remap[cls_ids]
921
+
922
+ keep = self._conf_filter_mask(scores, cls_ids)
923
+ boxes_xywh = boxes_xywh[keep]
924
+ scores = scores[keep]
925
+ cls_ids = cls_ids[keep]
926
+ if len(boxes_xywh) == 0:
927
+ return []
928
+ boxes = self._xywh_to_xyxy(boxes_xywh)
929
+
930
+ pad_w, pad_h = pad
931
+ boxes[:, [0, 2]] -= pad_w
932
+ boxes[:, [1, 3]] -= pad_h
933
+ boxes /= ratio
934
+ boxes = self._clip_boxes(boxes, orig_size)
935
+
936
+ boxes, scores, cls_ids = self._filter_sane_boxes(
937
+ boxes, scores, cls_ids, orig_size
938
+ )
939
+ if len(boxes) == 0:
940
+ return []
941
+
942
+ boxes, scores, cls_ids = self._per_view_pipeline(boxes, scores, cls_ids)
943
+ return self._build_results(boxes, scores, cls_ids)
944
+
945
+ def _postprocess(
946
+ self,
947
+ output: np.ndarray,
948
+ ratio: float,
949
+ pad: tuple[float, float],
950
+ orig_size: tuple[int, int],
951
+ ) -> list[BoundingBox]:
952
+ if output.ndim == 2 and output.shape[1] >= 6:
953
+ return self._decode_final_dets(output, ratio, pad, orig_size)
954
+ if output.ndim == 3 and output.shape[0] == 1 and output.shape[2] == 6:
955
+ return self._decode_final_dets(output, ratio, pad, orig_size)
956
+ return self._decode_raw_yolo(output, ratio, pad, orig_size)
957
+
958
+ def _predict_single(self, image: np.ndarray) -> list[BoundingBox]:
959
+ if image is None:
960
+ raise ValueError("Input image is None")
961
+ if not isinstance(image, np.ndarray):
962
+ raise TypeError(f"Input is not numpy array: {type(image)}")
963
+ if image.ndim != 3:
964
+ raise ValueError(f"Expected HWC image, got shape={image.shape}")
965
+ if image.shape[0] <= 0 or image.shape[1] <= 0:
966
+ raise ValueError(f"Invalid image shape={image.shape}")
967
+ if image.shape[2] != 3:
968
+ raise ValueError(f"Expected 3 channels, got shape={image.shape}")
969
+ if image.dtype != np.uint8:
970
+ image = image.astype(np.uint8)
971
+
972
+ input_tensor, ratio, pad, orig_size = self._preprocess(image)
973
+ expected = (1, 3, self.input_height, self.input_width)
974
+ if input_tensor.shape != expected:
975
+ raise ValueError(
976
+ f"Bad input tensor shape={input_tensor.shape}, expected={expected}"
977
+ )
978
+
979
+ outputs = self.session.run(self.output_names, {self.input_name: input_tensor})
980
+ return self._postprocess(outputs[0], ratio, pad, orig_size)
981
+
982
+ def _predict_tta(self, image: np.ndarray) -> list[BoundingBox]:
983
+ """Horizontal-flip TTA.
984
+
985
+ Strategy:
986
+ 1. Predict on original and on flipped image.
987
+ 2. Map flipped boxes back to original coordinates.
988
+ 3. Per-class hard NMS on the union.
989
+ 4. For each kept box, compute the max same-class score across the
990
+ FULL union (not just the post-NMS subset) -- this lets a high-
991
+ confidence flipped detection raise a borderline original one.
992
+ 5. Cross-class dedup to suppress same-physical-object multi-class.
993
+ 6. Smoke merge: overlapping / nested smoke boxes collapse into
994
+ their union (one box per smoke object).
995
+ """
996
+ boxes_orig = self._predict_single(image)
997
+ flipped = cv2.flip(image, 1)
998
+ boxes_flip = self._predict_single(flipped)
999
+ w = image.shape[1]
1000
+ boxes_flip = [
1001
+ BoundingBox(
1002
+ x1=w - b.x2, y1=b.y1, x2=w - b.x1, y2=b.y2,
1003
+ cls_id=b.cls_id, conf=b.conf,
1004
+ )
1005
+ for b in boxes_flip
1006
+ ]
1007
+ all_boxes = boxes_orig + boxes_flip
1008
+ if not all_boxes:
1009
+ return []
1010
+
1011
+ coords = np.array(
1012
+ [[b.x1, b.y1, b.x2, b.y2] for b in all_boxes], dtype=np.float32
1013
+ )
1014
+ scores = np.array([b.conf for b in all_boxes], dtype=np.float32)
1015
+ cls_ids = np.array([b.cls_id for b in all_boxes], dtype=np.int32)
1016
+ # view_id 0 = original, 1 = horizontal flip (mapped back to orig coords)
1017
+ view_ids = np.array(
1018
+ [0] * len(boxes_orig) + [1] * len(boxes_flip), dtype=np.int32
1019
+ )
1020
+
1021
+ hard_keep = self._per_class_hard_nms(coords, scores, cls_ids, self.iou_thres)
1022
+ if len(hard_keep) == 0:
1023
+ return []
1024
+ if len(hard_keep) > self.max_det:
1025
+ top = np.argsort(-scores[hard_keep])[: self.max_det]
1026
+ hard_keep = hard_keep[top]
1027
+
1028
+ boosted = self._max_score_per_cluster(
1029
+ coords[hard_keep], cls_ids[hard_keep],
1030
+ coords, scores, cls_ids, self.iou_thres,
1031
+ )
1032
+
1033
+ kept_coords = coords[hard_keep]
1034
+ kept_cls = cls_ids[hard_keep]
1035
+
1036
+ # Optional: drop low-conf detections seen in only one TTA view.
1037
+ if (
1038
+ self.use_tta_view_filter
1039
+ and self.tta_view_filter_max_conf > 0.0
1040
+ and len(kept_coords) > 0
1041
+ ):
1042
+ corrob = self._views_corroborated(
1043
+ kept_coords, kept_cls, coords, cls_ids, view_ids,
1044
+ self.tta_view_iou_thresh,
1045
+ )
1046
+ keep = ~((boosted <= self.tta_view_filter_max_conf) & (~corrob))
1047
+ kept_coords = kept_coords[keep]
1048
+ boosted = boosted[keep]
1049
+ kept_cls = kept_cls[keep]
1050
+
1051
+ if len(kept_coords) > 1:
1052
+ kept_coords, boosted, kept_cls = self._cross_class_dedup_op(
1053
+ kept_coords, boosted, kept_cls, self.cross_iou_thresh
1054
+ )
1055
+ if len(kept_coords) > 1:
1056
+ kept_coords, boosted, kept_cls = self._merge_same_class_boxes(
1057
+ kept_coords, boosted, kept_cls
1058
+ )
1059
+
1060
+ return [
1061
+ BoundingBox(
1062
+ x1=int(math.floor(kept_coords[j, 0])),
1063
+ y1=int(math.floor(kept_coords[j, 1])),
1064
+ x2=int(math.ceil(kept_coords[j, 2])),
1065
+ y2=int(math.ceil(kept_coords[j, 3])),
1066
+ cls_id=int(kept_cls[j]),
1067
+ conf=float(boosted[j]),
1068
+ )
1069
+ for j in range(len(kept_coords))
1070
+ ]
1071
+
1072
+ def predict_batch(
1073
+ self,
1074
+ batch_images: list[ndarray],
1075
+ offset: int,
1076
+ n_keypoints: int,
1077
+ ) -> list[TVFrameResult]:
1078
+ results: list[TVFrameResult] = []
1079
+ for frame_number_in_batch, image in enumerate(batch_images):
1080
+ try:
1081
+ if self.use_tta:
1082
+ boxes = self._predict_tta(image)
1083
+ else:
1084
+ boxes = self._predict_single(image)
1085
+ # Color-prior + edge FP filters on the merged result, in
1086
+ # original-image coords. Single insertion point so they run once
1087
+ # per frame for both the TTA and non-TTA paths.
1088
+ if isinstance(image, np.ndarray) and image.ndim == 3:
1089
+ boxes = self._filter_low_conf_by_color(image, boxes)
1090
+ boxes = self._remove_edge_low_conf(
1091
+ boxes, (image.shape[1], image.shape[0])
1092
+ )
1093
+ except Exception as e:
1094
+ print(
1095
+ f"⚠️ Inference failed for frame "
1096
+ f"{offset + frame_number_in_batch}: {e}"
1097
+ )
1098
+ boxes = []
1099
+ results.append(
1100
+ TVFrameResult(
1101
+ frame_id=offset + frame_number_in_batch,
1102
+ boxes=boxes,
1103
+ keypoints=[(0, 0) for _ in range(max(0, int(n_keypoints)))],
1104
+ )
1105
+ )
1106
+ return results
weights.onnx ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:922b2e3c6a96ae29371b263a7cca95c10c63c272f5ef703851d015401d60a68b
3
+ size 10459180