SuperBitDev commited on
Commit
54500ec
·
verified ·
1 Parent(s): d2064e5

Upload folder using huggingface_hub

Browse files
Files changed (2) hide show
  1. miner.py +293 -628
  2. weights.onnx +2 -2
miner.py CHANGED
@@ -1,8 +1,5 @@
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,293 +14,138 @@ class BoundingBox(BaseModel):
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,9 +166,7 @@ class Miner:
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,21 +183,13 @@ class Miner:
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,28 +201,14 @@ class Miner:
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,37 +223,20 @@ class Miner:
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,7 +244,7 @@ class Miner:
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,11 +253,8 @@ class Miner:
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,126 +262,61 @@ class Miner:
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,17 +326,12 @@ class Miner:
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,16 +349,9 @@ class Miner:
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,43 +368,30 @@ class Miner:
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,33 +403,25 @@ class Miner:
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,60 +429,34 @@ class Miner:
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,123 +466,170 @@ class Miner:
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,189 +639,133 @@ class Miner:
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
 
 
 
1
  from pathlib import Path
2
  import math
 
3
  import cv2
4
  import numpy as np
5
  import onnxruntime as ort
 
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.55
26
  cross_iou_thresh = 0.8
27
+ max_det = 30
28
+ _conf_thres_array = np.array([0.22, 0.30, 0.30], dtype=np.float32)
29
+ _bonus_array = np.array([0.05, 0.05, 0.05], 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.95
38
+ fire_expand = 1.05
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
  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
  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
  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
  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
  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
  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
  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
  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
  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
  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
  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
 
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
  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
  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
weights.onnx CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:0a157aa22ed68bb11e35546cd2ef2b3502d20ed5e2468c068544ad3c85298375
3
- size 9823990
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:9294fb5b3e873bb9239bc5728e97141c140999c827f1221b15bcd911eb756b74
3
+ size 9842050