MTerryJack commited on
Commit
744fdbb
·
verified ·
1 Parent(s): 41a5806

subnet_bridge: copy winning miner repo into library

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