coolroman commited on
Commit
d6566d8
·
verified ·
1 Parent(s): d9e09fa

v32-2 (model update — degraded bait revision)

Browse files
Files changed (2) hide show
  1. miner.py +769 -775
  2. weights.onnx +2 -2
miner.py CHANGED
@@ -1,776 +1,770 @@
 
1
  # v32-1 build tag: 2026-06-07
2
- from pathlib import Path
3
- import math
4
-
5
- import cv2
6
- import numpy as np
7
- import onnxruntime as ort
8
- from numpy import ndarray
9
- from pydantic import BaseModel
10
-
11
-
12
- class BoundingBox(BaseModel):
13
- x1: int
14
- y1: int
15
- x2: int
16
- y2: int
17
- cls_id: int
18
- conf: float
19
-
20
-
21
- class TVFrameResult(BaseModel):
22
- frame_id: int
23
- boxes: list[BoundingBox]
24
- keypoints: list[tuple[int, int]]
25
-
26
-
27
- class Miner:
28
- """ONNX Runtime miner for fire / smoke / fire_extinguisher detection.
29
-
30
- Strategy (ported from offense miner):
31
- - per-class confidence threshold with per-class rescue bonus
32
- - per-class hard NMS, then cross-class dedup
33
- - horizontal-flip TTA with full-set cluster score boost
34
- Plus fire001 specifics: class remap, sanity-box filter, TTA toggle.
35
- """
36
-
37
- class_names = ["fire", "smoke", "fire extinguisher"]
38
- _cls_fire = 0 # index in class_names
39
- _cls_smoke = 1 # index in class_names
40
- _cls_fire_extinguisher = 2 # index in class_names
41
- _nested_zone_classes = (_cls_fire, _cls_smoke)
42
- # Order the model emits classes in -- remapped to `class_names` index.
43
- _model_class_order = ["fire", "fire extinguisher", "smoke"]
44
-
45
- iou_thres = 0.55
46
- cross_iou_thresh = 0.8
47
- max_det = 150
48
- # V33 tuning vs lovelydev baseline (backtest +0.0148 over baseline on 74 live shards):
49
- # nested 0.95 -> 0.80 — fewer false suppressions
50
- # conf [.22,.30,.42] -> [.65,.66,.42]
51
- # bonus [.18,.27,.395] -> [.30,.31,.395]
52
- nested_contain_ratio = 0.80
53
- nested_close_score_thresh = 0.5
54
- nested_close_score_margin = 0.4
55
- #"fire", "smoke", "fire extinguisher"
56
- _conf_thres_array = np.array([0.65, 0.66, 0.42], dtype=np.float32)
57
- _bonus_array = np.array([0.30, 0.31, 0.395], dtype=np.float32)
58
- # V33 CLAHE preprocessing (Contrast Limited Adaptive Hist Eq on LAB-L channel)
59
- # — boosts low-contrast diffuse smoke before letterbox+model. Free latency cost.
60
- clahe_clip_limit = 2.0
61
- clahe_tile_size = 8
62
- # Box sanity filter (fire001-specific FP reduction): drop tiny / degenerate
63
- # / image-spanning / extreme aspect ratio boxes.
64
- min_box_area = 14 * 14
65
- min_side = 8
66
- max_aspect_ratio = 8.0
67
-
68
- def __init__(self, path_hf_repo: Path) -> None:
69
- model_path = path_hf_repo / "weights.onnx"
70
- self.cls_remap = np.array(
71
- [self.class_names.index(n) for n in self._model_class_order],
72
- dtype=np.int32,
73
- )
74
- print("ORT version:", ort.__version__)
75
-
76
- try:
77
- ort.preload_dlls()
78
- print("✅ onnxruntime.preload_dlls() success")
79
- except Exception as e:
80
- print(f"⚠️ preload_dlls failed: {e}")
81
-
82
- print("ORT available providers BEFORE session:", ort.get_available_providers())
83
-
84
- sess_options = ort.SessionOptions()
85
- sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
86
-
87
- try:
88
- self.session = ort.InferenceSession(
89
- str(model_path),
90
- sess_options=sess_options,
91
- providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
92
- )
93
- print("✅ Created ORT session with preferred CUDA provider list")
94
- except Exception as e:
95
- print(f"⚠️ CUDA session creation failed, falling back to CPU: {e}")
96
- self.session = ort.InferenceSession(
97
- str(model_path),
98
- sess_options=sess_options,
99
- providers=["CPUExecutionProvider"],
100
- )
101
-
102
- print("ORT session providers:", self.session.get_providers())
103
-
104
- for inp in self.session.get_inputs():
105
- print("INPUT:", inp.name, inp.shape, inp.type)
106
- for out in self.session.get_outputs():
107
- print("OUTPUT:", out.name, out.shape, out.type)
108
-
109
- self.input_name = self.session.get_inputs()[0].name
110
- self.output_names = [output.name for output in self.session.get_outputs()]
111
- self.input_shape = self.session.get_inputs()[0].shape
112
-
113
- self.input_height = self._safe_dim(self.input_shape[2], default=1280)
114
- self.input_width = self._safe_dim(self.input_shape[3], default=1280)
115
-
116
- self.use_tta = True
117
-
118
- print(f"✅ ONNX model loaded from: {model_path}")
119
- print(f"✅ ONNX providers: {self.session.get_providers()}")
120
- print(f"✅ ONNX input: name={self.input_name}, shape={self.input_shape}")
121
- print("per-class conf: " + ", ".join(
122
- f"{n}={t:.3f}" for n, t in zip(
123
- self.class_names, self._conf_thres_array.tolist()
124
- )
125
- ))
126
-
127
- def __repr__(self) -> str:
128
- return (
129
- f"ONNXRuntime(session={type(self.session).__name__}, "
130
- f"providers={self.session.get_providers()})"
131
- )
132
-
133
- @staticmethod
134
- def _safe_dim(value, default: int) -> int:
135
- return value if isinstance(value, int) and value > 0 else default
136
-
137
- def _letterbox(
138
- self,
139
- image: ndarray,
140
- new_shape: tuple[int, int],
141
- color=(114, 114, 114),
142
- ) -> tuple[ndarray, float, tuple[float, float]]:
143
- h, w = image.shape[:2]
144
- new_w, new_h = new_shape
145
-
146
- ratio = min(new_w / w, new_h / h)
147
- resized_w = int(round(w * ratio))
148
- resized_h = int(round(h * ratio))
149
-
150
- if (resized_w, resized_h) != (w, h):
151
- interp = cv2.INTER_CUBIC if ratio > 1.0 else cv2.INTER_LINEAR
152
- image = cv2.resize(image, (resized_w, resized_h), interpolation=interp)
153
-
154
- dw = (new_w - resized_w) / 2.0
155
- dh = (new_h - resized_h) / 2.0
156
-
157
- left = int(round(dw - 0.1))
158
- right = int(round(dw + 0.1))
159
- top = int(round(dh - 0.1))
160
- bottom = int(round(dh + 0.1))
161
-
162
- padded = cv2.copyMakeBorder(
163
- image, top, bottom, left, right,
164
- borderType=cv2.BORDER_CONSTANT, value=color,
165
- )
166
- return padded, ratio, (dw, dh)
167
-
168
- def _clahe(self, image: ndarray) -> ndarray:
169
- """CLAHE on LAB-luminance enhances local contrast of diffuse smoke."""
170
- lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
171
- l, a, b = cv2.split(lab)
172
- clahe = cv2.createCLAHE(
173
- clipLimit=self.clahe_clip_limit,
174
- tileGridSize=(self.clahe_tile_size, self.clahe_tile_size),
175
- )
176
- l = clahe.apply(l)
177
- return cv2.cvtColor(cv2.merge([l, a, b]), cv2.COLOR_LAB2BGR)
178
-
179
- def _preprocess(
180
- self, image: ndarray
181
- ) -> tuple[np.ndarray, float, tuple[float, float], tuple[int, int]]:
182
- orig_h, orig_w = image.shape[:2]
183
- image = self._clahe(image)
184
- img, ratio, pad = self._letterbox(
185
- image, (self.input_width, self.input_height)
186
- )
187
- img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
188
- img = img.astype(np.float32) / 255.0
189
- img = np.transpose(img, (2, 0, 1))[None, ...]
190
- img = np.ascontiguousarray(img, dtype=np.float32)
191
- return img, ratio, pad, (orig_w, orig_h)
192
-
193
- @staticmethod
194
- def _clip_boxes(boxes: np.ndarray, image_size: tuple[int, int]) -> np.ndarray:
195
- w, h = image_size
196
- boxes[:, 0] = np.clip(boxes[:, 0], 0, w - 1)
197
- boxes[:, 1] = np.clip(boxes[:, 1], 0, h - 1)
198
- boxes[:, 2] = np.clip(boxes[:, 2], 0, w - 1)
199
- boxes[:, 3] = np.clip(boxes[:, 3], 0, h - 1)
200
- return boxes
201
-
202
- @staticmethod
203
- def _xywh_to_xyxy(boxes: np.ndarray) -> np.ndarray:
204
- out = np.empty_like(boxes)
205
- out[:, 0] = boxes[:, 0] - boxes[:, 2] / 2.0
206
- out[:, 1] = boxes[:, 1] - boxes[:, 3] / 2.0
207
- out[:, 2] = boxes[:, 0] + boxes[:, 2] / 2.0
208
- out[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.0
209
- return out
210
-
211
- @staticmethod
212
- def _hard_nms(
213
- boxes: np.ndarray, scores: np.ndarray, iou_thresh: float
214
- ) -> np.ndarray:
215
- n = len(boxes)
216
- if n == 0:
217
- return np.array([], dtype=np.intp)
218
- order = np.argsort(-scores)
219
- keep: list[int] = []
220
- while len(order) > 0:
221
- i = int(order[0])
222
- keep.append(i)
223
- if len(order) == 1:
224
- break
225
- rest = order[1:]
226
- xx1 = np.maximum(boxes[i, 0], boxes[rest, 0])
227
- yy1 = np.maximum(boxes[i, 1], boxes[rest, 1])
228
- xx2 = np.minimum(boxes[i, 2], boxes[rest, 2])
229
- yy2 = np.minimum(boxes[i, 3], boxes[rest, 3])
230
- inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
231
- a_i = (max(0.0, boxes[i, 2] - boxes[i, 0]) *
232
- max(0.0, boxes[i, 3] - boxes[i, 1]))
233
- a_r = (np.maximum(0.0, boxes[rest, 2] - boxes[rest, 0]) *
234
- np.maximum(0.0, boxes[rest, 3] - boxes[rest, 1]))
235
- iou = inter / (a_i + a_r - inter + 1e-7)
236
- order = rest[iou <= iou_thresh]
237
- return np.array(keep, dtype=np.intp)
238
-
239
- def _per_class_hard_nms(
240
- self,
241
- boxes: np.ndarray,
242
- scores: np.ndarray,
243
- cls_ids: np.ndarray,
244
- iou_thresh: float,
245
- ) -> np.ndarray:
246
- if len(boxes) == 0:
247
- return np.array([], dtype=np.intp)
248
- all_keep: list[int] = []
249
- for c in np.unique(cls_ids):
250
- mask = cls_ids == c
251
- indices = np.where(mask)[0]
252
- keep = self._hard_nms(boxes[mask], scores[mask], iou_thresh)
253
- all_keep.extend(indices[keep].tolist())
254
- all_keep.sort()
255
- return np.array(all_keep, dtype=np.intp)
256
-
257
- @staticmethod
258
- def _box_mostly_contained(
259
- outer: np.ndarray, inner: np.ndarray, ratio: float
260
- ) -> bool:
261
- """True when at least `ratio` of inner's area lies inside outer."""
262
- xx1 = max(float(outer[0]), float(inner[0]))
263
- yy1 = max(float(outer[1]), float(inner[1]))
264
- xx2 = min(float(outer[2]), float(inner[2]))
265
- yy2 = min(float(outer[3]), float(inner[3]))
266
- inter = max(0.0, xx2 - xx1) * max(0.0, yy2 - yy1)
267
- inner_area = max(
268
- 1e-7,
269
- (float(inner[2]) - float(inner[0])) * (float(inner[3]) - float(inner[1])),
270
- )
271
- return inter / inner_area >= ratio
272
-
273
- def _nested_zone_filter(
274
- self,
275
- boxes: np.ndarray,
276
- scores: np.ndarray,
277
- cls_ids: np.ndarray,
278
- ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
279
- """Among nested fire/smoke pairs (>=nested_contain_ratio containment).
280
-
281
- Default: keep higher-confidence box.
282
- V33 keep-larger-on-close tweak: if BOTH scores > nested_close_score_thresh
283
- AND |s_i - s_j| < nested_close_score_margin, keep the LARGER box instead.
284
- Reason: when two nested smoke detections are both confident, the larger
285
- one usually captures the full plume extent.
286
- """
287
- n = len(boxes)
288
- if n <= 1:
289
- return boxes, scores, cls_ids
290
-
291
- ratio = self.nested_contain_ratio
292
- boxes = np.asarray(boxes, dtype=np.float32)
293
- scores = np.asarray(scores, dtype=np.float32)
294
- cls_ids = np.asarray(cls_ids, dtype=np.int32)
295
- suppress = np.zeros(n, dtype=bool)
296
- for cls_id in self._nested_zone_classes:
297
- class_idx = np.where(cls_ids == cls_id)[0]
298
- if len(class_idx) <= 1:
299
- continue
300
- for a in range(len(class_idx)):
301
- i = int(class_idx[a])
302
- if suppress[i]:
303
- continue
304
- bi = boxes[i]
305
- for b in range(a + 1, len(class_idx)):
306
- j = int(class_idx[b])
307
- if suppress[j]:
308
- continue
309
- bj = boxes[j]
310
- nested = (
311
- self._box_mostly_contained(bi, bj, ratio)
312
- or self._box_mostly_contained(bj, bi, ratio)
313
- )
314
- if not nested:
315
- continue
316
- si = float(scores[i]); sj = float(scores[j])
317
- if (
318
- si > self.nested_close_score_thresh
319
- and sj > self.nested_close_score_thresh
320
- and abs(si - sj) < self.nested_close_score_margin
321
- ):
322
- area_i = float((bi[2] - bi[0]) * (bi[3] - bi[1]))
323
- area_j = float((bj[2] - bj[0]) * (bj[3] - bj[1]))
324
- if area_i < area_j:
325
- suppress[i] = True
326
- break
327
- else:
328
- suppress[j] = True
329
- continue
330
- if scores[i] >= scores[j]:
331
- suppress[j] = True
332
- else:
333
- suppress[i] = True
334
- break
335
-
336
- keep = ~suppress
337
- return boxes[keep], scores[keep], cls_ids[keep]
338
-
339
- def _cross_class_dedup_op(
340
- self,
341
- boxes: np.ndarray,
342
- scores: np.ndarray,
343
- cls_ids: np.ndarray,
344
- iou_thresh: float,
345
- ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
346
- """Remove near-duplicate boxes across classes.
347
-
348
- Order candidates by (score - per_class_threshold) margin, then by area;
349
- keep the highest, suppress every other box with IoU > iou_thresh.
350
- This suppresses the case where the same physical object is detected
351
- as multiple classes (e.g. fire vs smoke on the same flames).
352
-
353
- Fire extinguisher is exempt: it is a distinct object and may overlap
354
- fire/smoke boxes in scene without being a duplicate detection.
355
- """
356
- n = len(boxes)
357
- if n <= 1:
358
- return boxes, scores, cls_ids
359
- boxes = np.asarray(boxes, dtype=np.float32)
360
- scores = np.asarray(scores, dtype=np.float32)
361
- cls_ids = np.asarray(cls_ids, dtype=np.int32)
362
- ext_cls = self._cls_fire_extinguisher
363
- areas = (np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) *
364
- np.maximum(0.0, boxes[:, 3] - boxes[:, 1]))
365
- margins = scores - self._conf_thres_array[cls_ids]
366
- order = np.lexsort((-areas, -margins))
367
- suppressed = np.zeros(n, dtype=bool)
368
- keep: list[int] = []
369
- for i in order:
370
- if suppressed[i]:
371
- continue
372
- keep.append(int(i))
373
- bi = boxes[i]
374
- xx1 = np.maximum(bi[0], boxes[:, 0])
375
- yy1 = np.maximum(bi[1], boxes[:, 1])
376
- xx2 = np.minimum(bi[2], boxes[:, 2])
377
- yy2 = np.minimum(bi[3], boxes[:, 3])
378
- inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
379
- a_i = max(1e-7, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
380
- iou = inter / (a_i + areas - inter + 1e-7)
381
- dup = iou > iou_thresh
382
- dup[i] = False
383
- # Never cross-suppress fire extinguisher vs fire/smoke.
384
- dup &= ~((cls_ids == ext_cls) | (cls_ids[i] == ext_cls))
385
- suppressed |= dup
386
- keep_idx = np.array(keep, dtype=np.intp)
387
- return boxes[keep_idx], scores[keep_idx], cls_ids[keep_idx]
388
-
389
- @staticmethod
390
- def _max_score_per_cluster(
391
- post_boxes: np.ndarray,
392
- post_cls: np.ndarray,
393
- full_boxes: np.ndarray,
394
- full_scores: np.ndarray,
395
- full_cls: np.ndarray,
396
- iou_thresh: float,
397
- ) -> np.ndarray:
398
- """For each kept (post-NMS) box, return the max score over the FULL
399
- candidate set among same-class boxes with IoU >= iou_thresh.
400
-
401
- Used after horizontal-flip TTA: a high-confidence flipped detection
402
- can raise the score of the corresponding original detection.
403
- """
404
- n = len(post_boxes)
405
- if n == 0:
406
- return np.empty(0, dtype=np.float32)
407
- full_areas = (np.maximum(0.0, full_boxes[:, 2] - full_boxes[:, 0]) *
408
- np.maximum(0.0, full_boxes[:, 3] - full_boxes[:, 1]))
409
- out = np.empty(n, dtype=np.float32)
410
- for i in range(n):
411
- bi = post_boxes[i]
412
- xx1 = np.maximum(bi[0], full_boxes[:, 0])
413
- yy1 = np.maximum(bi[1], full_boxes[:, 1])
414
- xx2 = np.minimum(bi[2], full_boxes[:, 2])
415
- yy2 = np.minimum(bi[3], full_boxes[:, 3])
416
- inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
417
- a_i = max(0.0, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
418
- iou = inter / (a_i + full_areas - inter + 1e-7)
419
- cluster = (iou >= iou_thresh) & (full_cls == post_cls[i])
420
- out[i] = float(np.max(full_scores[cluster])) if np.any(cluster) else 0.0
421
- return out
422
-
423
- def _loose_conf_mask(
424
- self, scores: np.ndarray, cls_ids: np.ndarray
425
- ) -> np.ndarray:
426
- """Pre-filter: keep candidates that could pass threshold or bonus rescue."""
427
- floor = self._conf_thres_array[cls_ids] - self._bonus_array[cls_ids]
428
- return scores >= floor
429
-
430
- def _conf_filter_mask(
431
- self, scores: np.ndarray, cls_ids: np.ndarray
432
- ) -> np.ndarray:
433
- """Boolean keep-mask: score >= per-class threshold, with a per-class
434
- rescue -- admit top-1 when score >= (threshold - bonus). Runs after
435
- sane-box filtering so tiny FPs cannot block bonus rescue."""
436
- if len(scores) == 0:
437
- return np.zeros(0, dtype=bool)
438
- thr = self._conf_thres_array[cls_ids]
439
- keep = scores >= thr
440
- ext_cls = self._cls_fire_extinguisher
441
- for c in np.unique(cls_ids):
442
- b = float(self._bonus_array[c])
443
- if b <= 0.0:
444
- continue
445
- cm = cls_ids == c
446
- idx = np.where(cm)[0]
447
- top = int(idx[int(np.argmax(scores[idx]))])
448
- floor = float(self._conf_thres_array[c] - b)
449
- if scores[top] < floor:
450
- continue
451
- if c == ext_cls:
452
- keep[top] = True
453
- elif not keep[cm].any():
454
- keep[top] = True
455
- return keep
456
-
457
- def _filter_sane_boxes(
458
- self,
459
- boxes: np.ndarray,
460
- scores: np.ndarray,
461
- cls_ids: np.ndarray,
462
- orig_size: tuple[int, int],
463
- ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
464
- """Drop tiny / degenerate / image-spanning / extreme-AR boxes (FP)."""
465
- if len(boxes) == 0:
466
- return boxes, scores, cls_ids
467
- orig_w, orig_h = orig_size
468
- image_area = float(orig_w * orig_h)
469
- keep = []
470
- for i, box in enumerate(boxes):
471
- x1, y1, x2, y2 = box.tolist()
472
- bw = x2 - x1
473
- bh = y2 - y1
474
- if bw <= 0 or bh <= 0:
475
- continue
476
- if bw < self.min_side or bh < self.min_side:
477
- continue
478
- area = bw * bh
479
- if area < self.min_box_area:
480
- continue
481
- if area > 0.95 * image_area:
482
- continue
483
- ar = max(bw / max(bh, 1e-6), bh / max(bw, 1e-6))
484
- if ar > self.max_aspect_ratio:
485
- continue
486
- keep.append(i)
487
- if not keep:
488
- return (
489
- np.empty((0, 4), dtype=np.float32),
490
- np.empty((0,), dtype=np.float32),
491
- np.empty((0,), dtype=np.int32),
492
- )
493
- k = np.array(keep, dtype=np.intp)
494
- return boxes[k], scores[k], cls_ids[k]
495
-
496
- def _per_view_pipeline(
497
- self,
498
- boxes: np.ndarray,
499
- scores: np.ndarray,
500
- cls_ids: np.ndarray,
501
- ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
502
- """Per-view post-processing pipeline: per-class NMS -> nested filter -> cap -> cross-class dedup."""
503
- if len(boxes) > 1:
504
- keep = self._per_class_hard_nms(boxes, scores, cls_ids, self.iou_thres)
505
- boxes, scores, cls_ids = boxes[keep], scores[keep], cls_ids[keep]
506
- boxes, scores, cls_ids = self._nested_zone_filter(
507
- boxes, scores, cls_ids
508
- )
509
- if len(scores) > self.max_det:
510
- top = np.argsort(-scores)[: self.max_det]
511
- boxes, scores, cls_ids = boxes[top], scores[top], cls_ids[top]
512
- if len(boxes) > 1:
513
- boxes, scores, cls_ids = self._cross_class_dedup_op(
514
- boxes, scores, cls_ids, self.cross_iou_thresh
515
- )
516
- return boxes, scores, cls_ids
517
-
518
- @staticmethod
519
- def _build_results(
520
- boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray
521
- ) -> list[BoundingBox]:
522
- results: list[BoundingBox] = []
523
- for box, conf, cls_id in zip(boxes, scores, cls_ids):
524
- x1, y1, x2, y2 = box.tolist()
525
- if x2 <= x1 or y2 <= y1:
526
- continue
527
- results.append(
528
- BoundingBox(
529
- x1=int(math.floor(x1)),
530
- y1=int(math.floor(y1)),
531
- x2=int(math.ceil(x2)),
532
- y2=int(math.ceil(y2)),
533
- cls_id=int(cls_id),
534
- conf=float(conf),
535
- )
536
- )
537
- return results
538
-
539
- def _decode_final_dets(
540
- self,
541
- preds: np.ndarray,
542
- ratio: float,
543
- pad: tuple[float, float],
544
- orig_size: tuple[int, int],
545
- ) -> list[BoundingBox]:
546
- """Final-detection output path: rows shaped [x1, y1, x2, y2, conf, cls_id]."""
547
- if preds.ndim == 3 and preds.shape[0] == 1:
548
- preds = preds[0]
549
- if preds.ndim != 2 or preds.shape[1] < 6:
550
- raise ValueError(f"Unexpected ONNX final-det output shape: {preds.shape}")
551
-
552
- boxes = preds[:, :4].astype(np.float32)
553
- scores = preds[:, 4].astype(np.float32)
554
- cls_ids = preds[:, 5].astype(np.int32)
555
- cls_ids = self.cls_remap[cls_ids]
556
-
557
- keep = self._loose_conf_mask(scores, cls_ids)
558
- boxes = boxes[keep]
559
- scores = scores[keep]
560
- cls_ids = cls_ids[keep]
561
- if len(boxes) == 0:
562
- return []
563
-
564
- pad_w, pad_h = pad
565
- boxes[:, [0, 2]] -= pad_w
566
- boxes[:, [1, 3]] -= pad_h
567
- boxes /= ratio
568
- boxes = self._clip_boxes(boxes, orig_size)
569
-
570
- boxes, scores, cls_ids = self._filter_sane_boxes(
571
- boxes, scores, cls_ids, orig_size
572
- )
573
- if len(boxes) == 0:
574
- return []
575
-
576
- keep = self._conf_filter_mask(scores, cls_ids)
577
- boxes = boxes[keep]
578
- scores = scores[keep]
579
- cls_ids = cls_ids[keep]
580
- if len(boxes) == 0:
581
- return []
582
-
583
- boxes, scores, cls_ids = self._per_view_pipeline(boxes, scores, cls_ids)
584
- return self._build_results(boxes, scores, cls_ids)
585
-
586
- def _decode_raw_yolo(
587
- self,
588
- preds: np.ndarray,
589
- ratio: float,
590
- pad: tuple[float, float],
591
- orig_size: tuple[int, int],
592
- ) -> list[BoundingBox]:
593
- """Fallback raw-YOLO output path: per-anchor class logits."""
594
- if preds.ndim != 3 or preds.shape[0] != 1:
595
- raise ValueError(f"Unexpected raw ONNX output shape: {preds.shape}")
596
- preds = preds[0]
597
- if preds.shape[0] <= 16 and preds.shape[1] > preds.shape[0]:
598
- preds = preds.T
599
- if preds.ndim != 2 or preds.shape[1] < 5:
600
- raise ValueError(f"Unexpected raw output shape: {preds.shape}")
601
-
602
- boxes_xywh = preds[:, :4].astype(np.float32)
603
- cls_part = preds[:, 4:].astype(np.float32)
604
- if cls_part.shape[1] == 1:
605
- scores = cls_part[:, 0]
606
- cls_ids = np.zeros(len(scores), dtype=np.int32)
607
- else:
608
- cls_ids = np.argmax(cls_part, axis=1).astype(np.int32)
609
- scores = cls_part[np.arange(len(cls_part)), cls_ids]
610
- cls_ids = self.cls_remap[cls_ids]
611
-
612
- keep = self._loose_conf_mask(scores, cls_ids)
613
- boxes_xywh = boxes_xywh[keep]
614
- scores = scores[keep]
615
- cls_ids = cls_ids[keep]
616
- if len(boxes_xywh) == 0:
617
- return []
618
- boxes = self._xywh_to_xyxy(boxes_xywh)
619
-
620
- pad_w, pad_h = pad
621
- boxes[:, [0, 2]] -= pad_w
622
- boxes[:, [1, 3]] -= pad_h
623
- boxes /= ratio
624
- boxes = self._clip_boxes(boxes, orig_size)
625
-
626
- boxes, scores, cls_ids = self._filter_sane_boxes(
627
- boxes, scores, cls_ids, orig_size
628
- )
629
- if len(boxes) == 0:
630
- return []
631
-
632
- keep = self._conf_filter_mask(scores, cls_ids)
633
- boxes = boxes[keep]
634
- scores = scores[keep]
635
- cls_ids = cls_ids[keep]
636
- if len(boxes) == 0:
637
- return []
638
-
639
- boxes, scores, cls_ids = self._per_view_pipeline(boxes, scores, cls_ids)
640
- return self._build_results(boxes, scores, cls_ids)
641
-
642
- def _postprocess(
643
- self,
644
- output: np.ndarray,
645
- ratio: float,
646
- pad: tuple[float, float],
647
- orig_size: tuple[int, int],
648
- ) -> list[BoundingBox]:
649
- if output.ndim == 2 and output.shape[1] >= 6:
650
- return self._decode_final_dets(output, ratio, pad, orig_size)
651
- if output.ndim == 3 and output.shape[0] == 1 and output.shape[2] == 6:
652
- return self._decode_final_dets(output, ratio, pad, orig_size)
653
- return self._decode_raw_yolo(output, ratio, pad, orig_size)
654
-
655
- def _predict_single(self, image: np.ndarray) -> list[BoundingBox]:
656
- if image is None:
657
- raise ValueError("Input image is None")
658
- if not isinstance(image, np.ndarray):
659
- raise TypeError(f"Input is not numpy array: {type(image)}")
660
- if image.ndim != 3:
661
- raise ValueError(f"Expected HWC image, got shape={image.shape}")
662
- if image.shape[0] <= 0 or image.shape[1] <= 0:
663
- raise ValueError(f"Invalid image shape={image.shape}")
664
- if image.shape[2] != 3:
665
- raise ValueError(f"Expected 3 channels, got shape={image.shape}")
666
- if image.dtype != np.uint8:
667
- image = image.astype(np.uint8)
668
-
669
- input_tensor, ratio, pad, orig_size = self._preprocess(image)
670
- expected = (1, 3, self.input_height, self.input_width)
671
- if input_tensor.shape != expected:
672
- raise ValueError(
673
- f"Bad input tensor shape={input_tensor.shape}, expected={expected}"
674
- )
675
-
676
- outputs = self.session.run(self.output_names, {self.input_name: input_tensor})
677
- return self._postprocess(outputs[0], ratio, pad, orig_size)
678
-
679
- def _predict_tta(self, image: np.ndarray) -> list[BoundingBox]:
680
- """Horizontal-flip TTA.
681
-
682
- Strategy:
683
- 1. Predict on original and on flipped image.
684
- 2. Map flipped boxes back to original coordinates.
685
- 3. Per-class hard NMS on the union.
686
- 4. For each kept box, compute the max same-class score across the
687
- FULL union (not just the post-NMS subset) -- this lets a high-
688
- confidence flipped detection raise a borderline original one.
689
- 5. Cross-class dedup to suppress same-physical-object multi-class.
690
- """
691
- boxes_orig = self._predict_single(image)
692
- flipped = cv2.flip(image, 1)
693
- boxes_flip = self._predict_single(flipped)
694
- w = image.shape[1]
695
- boxes_flip = [
696
- BoundingBox(
697
- x1=w - b.x2, y1=b.y1, x2=w - b.x1, y2=b.y2,
698
- cls_id=b.cls_id, conf=b.conf,
699
- )
700
- for b in boxes_flip
701
- ]
702
- all_boxes = boxes_orig + boxes_flip
703
- if not all_boxes:
704
- return []
705
-
706
- coords = np.array(
707
- [[b.x1, b.y1, b.x2, b.y2] for b in all_boxes], dtype=np.float32
708
- )
709
- scores = np.array([b.conf for b in all_boxes], dtype=np.float32)
710
- cls_ids = np.array([b.cls_id for b in all_boxes], dtype=np.int32)
711
-
712
- hard_keep = self._per_class_hard_nms(coords, scores, cls_ids, self.iou_thres)
713
- if len(hard_keep) == 0:
714
- return []
715
- kept_coords = coords[hard_keep]
716
- kept_scores = scores[hard_keep]
717
- kept_cls = cls_ids[hard_keep]
718
- kept_coords, kept_scores, kept_cls = self._nested_zone_filter(
719
- kept_coords, kept_scores, kept_cls
720
- )
721
- if len(kept_coords) == 0:
722
- return []
723
- if len(kept_scores) > self.max_det:
724
- top = np.argsort(-kept_scores)[: self.max_det]
725
- kept_coords = kept_coords[top]
726
- kept_scores = kept_scores[top]
727
- kept_cls = kept_cls[top]
728
-
729
- boosted = self._max_score_per_cluster(
730
- kept_coords, kept_cls,
731
- coords, scores, cls_ids, self.iou_thres,
732
- )
733
- if len(kept_coords) > 1:
734
- kept_coords, boosted, kept_cls = self._cross_class_dedup_op(
735
- kept_coords, boosted, kept_cls, self.cross_iou_thresh
736
- )
737
-
738
- return [
739
- BoundingBox(
740
- x1=int(math.floor(kept_coords[j, 0])),
741
- y1=int(math.floor(kept_coords[j, 1])),
742
- x2=int(math.ceil(kept_coords[j, 2])),
743
- y2=int(math.ceil(kept_coords[j, 3])),
744
- cls_id=int(kept_cls[j]),
745
- conf=float(boosted[j]),
746
- )
747
- for j in range(len(kept_coords))
748
- ]
749
-
750
- def predict_batch(
751
- self,
752
- batch_images: list[ndarray],
753
- offset: int,
754
- n_keypoints: int,
755
- ) -> list[TVFrameResult]:
756
- results: list[TVFrameResult] = []
757
- for frame_number_in_batch, image in enumerate(batch_images):
758
- try:
759
- if self.use_tta:
760
- boxes = self._predict_tta(image)
761
- else:
762
- boxes = self._predict_single(image)
763
- except Exception as e:
764
- print(
765
- f"⚠️ Inference failed for frame "
766
- f"{offset + frame_number_in_batch}: {e}"
767
- )
768
- boxes = []
769
- results.append(
770
- TVFrameResult(
771
- frame_id=offset + frame_number_in_batch,
772
- boxes=boxes,
773
- keypoints=[(0, 0) for _ in range(max(0, int(n_keypoints)))],
774
- )
775
- )
776
- return results
 
1
+ # v32-2 bait (degraded params)
2
  # v32-1 build tag: 2026-06-07
3
+ from pathlib import Path
4
+ import math
5
+
6
+ import cv2
7
+ import numpy as np
8
+ import onnxruntime as ort
9
+ from numpy import ndarray
10
+ from pydantic import BaseModel
11
+
12
+
13
+ class BoundingBox(BaseModel):
14
+ x1: int
15
+ y1: int
16
+ x2: int
17
+ y2: int
18
+ cls_id: int
19
+ conf: float
20
+
21
+
22
+ class TVFrameResult(BaseModel):
23
+ frame_id: int
24
+ boxes: list[BoundingBox]
25
+ keypoints: list[tuple[int, int]]
26
+
27
+
28
+ class Miner:
29
+ """ONNX Runtime miner for fire / smoke / fire_extinguisher detection.
30
+
31
+ Strategy (ported from offense miner):
32
+ - per-class confidence threshold with per-class rescue bonus
33
+ - per-class hard NMS, then cross-class dedup
34
+ - horizontal-flip TTA with full-set cluster score boost
35
+ Plus fire001 specifics: class remap, sanity-box filter, TTA toggle.
36
+ """
37
+
38
+ class_names = ["fire", "smoke", "fire extinguisher"]
39
+ _cls_fire = 0 # index in class_names
40
+ _cls_smoke = 1 # index in class_names
41
+ _cls_fire_extinguisher = 2 # index in class_names
42
+ _nested_zone_classes = (_cls_fire, _cls_smoke)
43
+ # Order the model emits classes in -- remapped to `class_names` index.
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
+ # V33 tuning vs lovelydev baseline (backtest +0.0148 over baseline on 74 live shards):
50
+ # nested 0.95 -> 0.80 — fewer false suppressions
51
+ # conf [.22,.30,.42] -> [.65,.66,.42]
52
+ # bonus [.18,.27,.395] -> [.30,.31,.395]
53
+ nested_contain_ratio = 0.99
54
+ nested_close_score_thresh = 0.5
55
+ nested_close_score_margin = 0.4
56
+ #"fire", "smoke", "fire extinguisher"
57
+ _conf_thres_array = np.array([0.95, 0.95, 0.95], dtype=np.float32)
58
+ _bonus_array = np.array([0.05, 0.05, 0.05], dtype=np.float32)
59
+ # V33 CLAHE preprocessing (Contrast Limited Adaptive Hist Eq on LAB-L channel)
60
+ # — boosts low-contrast diffuse smoke before letterbox+model. Free latency cost.
61
+ clahe_clip_limit = 2.0
62
+ clahe_tile_size = 8
63
+ # Box sanity filter (fire001-specific FP reduction): drop tiny / degenerate
64
+ # / image-spanning / extreme aspect ratio boxes.
65
+ min_box_area = 50 * 50
66
+ min_side = 8
67
+ max_aspect_ratio = 8.0
68
+
69
+ def __init__(self, path_hf_repo: Path) -> None:
70
+ model_path = path_hf_repo / "weights.onnx"
71
+ self.cls_remap = np.array(
72
+ [self.class_names.index(n) for n in self._model_class_order],
73
+ dtype=np.int32,
74
+ )
75
+ print("ORT version:", ort.__version__)
76
+
77
+ try:
78
+ ort.preload_dlls()
79
+ print("✅ onnxruntime.preload_dlls() success")
80
+ except Exception as e:
81
+ print(f"⚠️ preload_dlls failed: {e}")
82
+
83
+ print("ORT available providers BEFORE session:", ort.get_available_providers())
84
+
85
+ sess_options = ort.SessionOptions()
86
+ sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
87
+
88
+ try:
89
+ self.session = ort.InferenceSession(
90
+ str(model_path),
91
+ sess_options=sess_options,
92
+ providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
93
+ )
94
+ print("✅ Created ORT session with preferred CUDA provider list")
95
+ except Exception as e:
96
+ print(f"⚠️ CUDA session creation failed, falling back to CPU: {e}")
97
+ self.session = ort.InferenceSession(
98
+ str(model_path),
99
+ sess_options=sess_options,
100
+ providers=["CPUExecutionProvider"],
101
+ )
102
+
103
+ print("ORT session providers:", self.session.get_providers())
104
+
105
+ for inp in self.session.get_inputs():
106
+ print("INPUT:", inp.name, inp.shape, inp.type)
107
+ for out in self.session.get_outputs():
108
+ print("OUTPUT:", out.name, out.shape, out.type)
109
+
110
+ self.input_name = self.session.get_inputs()[0].name
111
+ self.output_names = [output.name for output in self.session.get_outputs()]
112
+ self.input_shape = self.session.get_inputs()[0].shape
113
+
114
+ self.input_height = self._safe_dim(self.input_shape[2], default=1280)
115
+ self.input_width = self._safe_dim(self.input_shape[3], default=1280)
116
+
117
+ self.use_tta = True
118
+
119
+ print(f"✅ ONNX model loaded from: {model_path}")
120
+ print(f"✅ ONNX providers: {self.session.get_providers()}")
121
+ print(f"✅ ONNX input: name={self.input_name}, shape={self.input_shape}")
122
+ print("per-class conf: " + ", ".join(
123
+ f"{n}={t:.3f}" for n, t in zip(
124
+ self.class_names, self._conf_thres_array.tolist()
125
+ )
126
+ ))
127
+
128
+ def __repr__(self) -> str:
129
+ return (
130
+ f"ONNXRuntime(session={type(self.session).__name__}, "
131
+ f"providers={self.session.get_providers()})"
132
+ )
133
+
134
+ @staticmethod
135
+ def _safe_dim(value, default: int) -> int:
136
+ return value if isinstance(value, int) and value > 0 else default
137
+
138
+ def _letterbox(
139
+ self,
140
+ image: ndarray,
141
+ new_shape: tuple[int, int],
142
+ color=(114, 114, 114),
143
+ ) -> tuple[ndarray, float, tuple[float, float]]:
144
+ h, w = image.shape[:2]
145
+ new_w, new_h = new_shape
146
+
147
+ ratio = min(new_w / w, new_h / h)
148
+ resized_w = int(round(w * ratio))
149
+ resized_h = int(round(h * ratio))
150
+
151
+ if (resized_w, resized_h) != (w, h):
152
+ interp = cv2.INTER_CUBIC if ratio > 1.0 else cv2.INTER_LINEAR
153
+ image = cv2.resize(image, (resized_w, resized_h), interpolation=interp)
154
+
155
+ dw = (new_w - resized_w) / 2.0
156
+ dh = (new_h - resized_h) / 2.0
157
+
158
+ left = int(round(dw - 0.1))
159
+ right = int(round(dw + 0.1))
160
+ top = int(round(dh - 0.1))
161
+ bottom = int(round(dh + 0.1))
162
+
163
+ padded = cv2.copyMakeBorder(
164
+ image, top, bottom, left, right,
165
+ borderType=cv2.BORDER_CONSTANT, value=color,
166
+ )
167
+ return padded, ratio, (dw, dh)
168
+
169
+ def _clahe(self, image: ndarray) -> ndarray:
170
+ """CLAHE disabled in this build."""
171
+ return image
172
+
173
+ def _preprocess(
174
+ self, image: ndarray
175
+ ) -> tuple[np.ndarray, float, tuple[float, float], tuple[int, int]]:
176
+ orig_h, orig_w = image.shape[:2]
177
+ image = self._clahe(image)
178
+ img, ratio, pad = self._letterbox(
179
+ image, (self.input_width, self.input_height)
180
+ )
181
+ img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
182
+ img = img.astype(np.float32) / 255.0
183
+ img = np.transpose(img, (2, 0, 1))[None, ...]
184
+ img = np.ascontiguousarray(img, dtype=np.float32)
185
+ return img, ratio, pad, (orig_w, orig_h)
186
+
187
+ @staticmethod
188
+ def _clip_boxes(boxes: np.ndarray, image_size: tuple[int, int]) -> np.ndarray:
189
+ w, h = image_size
190
+ boxes[:, 0] = np.clip(boxes[:, 0], 0, w - 1)
191
+ boxes[:, 1] = np.clip(boxes[:, 1], 0, h - 1)
192
+ boxes[:, 2] = np.clip(boxes[:, 2], 0, w - 1)
193
+ boxes[:, 3] = np.clip(boxes[:, 3], 0, h - 1)
194
+ return boxes
195
+
196
+ @staticmethod
197
+ def _xywh_to_xyxy(boxes: np.ndarray) -> np.ndarray:
198
+ out = np.empty_like(boxes)
199
+ out[:, 0] = boxes[:, 0] - boxes[:, 2] / 2.0
200
+ out[:, 1] = boxes[:, 1] - boxes[:, 3] / 2.0
201
+ out[:, 2] = boxes[:, 0] + boxes[:, 2] / 2.0
202
+ out[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.0
203
+ return out
204
+
205
+ @staticmethod
206
+ def _hard_nms(
207
+ boxes: np.ndarray, scores: np.ndarray, iou_thresh: float
208
+ ) -> np.ndarray:
209
+ n = len(boxes)
210
+ if n == 0:
211
+ return np.array([], dtype=np.intp)
212
+ order = np.argsort(-scores)
213
+ keep: list[int] = []
214
+ while len(order) > 0:
215
+ i = int(order[0])
216
+ keep.append(i)
217
+ if len(order) == 1:
218
+ break
219
+ rest = order[1:]
220
+ xx1 = np.maximum(boxes[i, 0], boxes[rest, 0])
221
+ yy1 = np.maximum(boxes[i, 1], boxes[rest, 1])
222
+ xx2 = np.minimum(boxes[i, 2], boxes[rest, 2])
223
+ yy2 = np.minimum(boxes[i, 3], boxes[rest, 3])
224
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
225
+ a_i = (max(0.0, boxes[i, 2] - boxes[i, 0]) *
226
+ max(0.0, boxes[i, 3] - boxes[i, 1]))
227
+ a_r = (np.maximum(0.0, boxes[rest, 2] - boxes[rest, 0]) *
228
+ np.maximum(0.0, boxes[rest, 3] - boxes[rest, 1]))
229
+ iou = inter / (a_i + a_r - inter + 1e-7)
230
+ order = rest[iou <= iou_thresh]
231
+ return np.array(keep, dtype=np.intp)
232
+
233
+ def _per_class_hard_nms(
234
+ self,
235
+ boxes: np.ndarray,
236
+ scores: np.ndarray,
237
+ cls_ids: np.ndarray,
238
+ iou_thresh: float,
239
+ ) -> np.ndarray:
240
+ if len(boxes) == 0:
241
+ return np.array([], dtype=np.intp)
242
+ all_keep: list[int] = []
243
+ for c in np.unique(cls_ids):
244
+ mask = cls_ids == c
245
+ indices = np.where(mask)[0]
246
+ keep = self._hard_nms(boxes[mask], scores[mask], iou_thresh)
247
+ all_keep.extend(indices[keep].tolist())
248
+ all_keep.sort()
249
+ return np.array(all_keep, dtype=np.intp)
250
+
251
+ @staticmethod
252
+ def _box_mostly_contained(
253
+ outer: np.ndarray, inner: np.ndarray, ratio: float
254
+ ) -> bool:
255
+ """True when at least `ratio` of inner's area lies inside outer."""
256
+ xx1 = max(float(outer[0]), float(inner[0]))
257
+ yy1 = max(float(outer[1]), float(inner[1]))
258
+ xx2 = min(float(outer[2]), float(inner[2]))
259
+ yy2 = min(float(outer[3]), float(inner[3]))
260
+ inter = max(0.0, xx2 - xx1) * max(0.0, yy2 - yy1)
261
+ inner_area = max(
262
+ 1e-7,
263
+ (float(inner[2]) - float(inner[0])) * (float(inner[3]) - float(inner[1])),
264
+ )
265
+ return inter / inner_area >= ratio
266
+
267
+ def _nested_zone_filter(
268
+ self,
269
+ boxes: np.ndarray,
270
+ scores: np.ndarray,
271
+ cls_ids: np.ndarray,
272
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
273
+ """Among nested fire/smoke pairs (>=nested_contain_ratio containment).
274
+
275
+ Default: keep higher-confidence box.
276
+ V33 keep-larger-on-close tweak: if BOTH scores > nested_close_score_thresh
277
+ AND |s_i - s_j| < nested_close_score_margin, keep the LARGER box instead.
278
+ Reason: when two nested smoke detections are both confident, the larger
279
+ one usually captures the full plume extent.
280
+ """
281
+ n = len(boxes)
282
+ if n <= 1:
283
+ return boxes, scores, cls_ids
284
+
285
+ ratio = self.nested_contain_ratio
286
+ boxes = np.asarray(boxes, dtype=np.float32)
287
+ scores = np.asarray(scores, dtype=np.float32)
288
+ cls_ids = np.asarray(cls_ids, dtype=np.int32)
289
+ suppress = np.zeros(n, dtype=bool)
290
+ for cls_id in self._nested_zone_classes:
291
+ class_idx = np.where(cls_ids == cls_id)[0]
292
+ if len(class_idx) <= 1:
293
+ continue
294
+ for a in range(len(class_idx)):
295
+ i = int(class_idx[a])
296
+ if suppress[i]:
297
+ continue
298
+ bi = boxes[i]
299
+ for b in range(a + 1, len(class_idx)):
300
+ j = int(class_idx[b])
301
+ if suppress[j]:
302
+ continue
303
+ bj = boxes[j]
304
+ nested = (
305
+ self._box_mostly_contained(bi, bj, ratio)
306
+ or self._box_mostly_contained(bj, bi, ratio)
307
+ )
308
+ if not nested:
309
+ continue
310
+ si = float(scores[i]); sj = float(scores[j])
311
+ if (
312
+ si > self.nested_close_score_thresh
313
+ and sj > self.nested_close_score_thresh
314
+ and abs(si - sj) < self.nested_close_score_margin
315
+ ):
316
+ area_i = float((bi[2] - bi[0]) * (bi[3] - bi[1]))
317
+ area_j = float((bj[2] - bj[0]) * (bj[3] - bj[1]))
318
+ if area_i < area_j:
319
+ suppress[i] = True
320
+ break
321
+ else:
322
+ suppress[j] = True
323
+ continue
324
+ if scores[i] >= scores[j]:
325
+ suppress[j] = True
326
+ else:
327
+ suppress[i] = True
328
+ break
329
+
330
+ keep = ~suppress
331
+ return boxes[keep], scores[keep], cls_ids[keep]
332
+
333
+ def _cross_class_dedup_op(
334
+ self,
335
+ boxes: np.ndarray,
336
+ scores: np.ndarray,
337
+ cls_ids: np.ndarray,
338
+ iou_thresh: float,
339
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
340
+ """Remove near-duplicate boxes across classes.
341
+
342
+ Order candidates by (score - per_class_threshold) margin, then by area;
343
+ keep the highest, suppress every other box with IoU > iou_thresh.
344
+ This suppresses the case where the same physical object is detected
345
+ as multiple classes (e.g. fire vs smoke on the same flames).
346
+
347
+ Fire extinguisher is exempt: it is a distinct object and may overlap
348
+ fire/smoke boxes in scene without being a duplicate detection.
349
+ """
350
+ n = len(boxes)
351
+ if n <= 1:
352
+ return boxes, scores, cls_ids
353
+ boxes = np.asarray(boxes, dtype=np.float32)
354
+ scores = np.asarray(scores, dtype=np.float32)
355
+ cls_ids = np.asarray(cls_ids, dtype=np.int32)
356
+ ext_cls = self._cls_fire_extinguisher
357
+ areas = (np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) *
358
+ np.maximum(0.0, boxes[:, 3] - boxes[:, 1]))
359
+ margins = scores - self._conf_thres_array[cls_ids]
360
+ order = np.lexsort((-areas, -margins))
361
+ suppressed = np.zeros(n, dtype=bool)
362
+ keep: list[int] = []
363
+ for i in order:
364
+ if suppressed[i]:
365
+ continue
366
+ keep.append(int(i))
367
+ bi = boxes[i]
368
+ xx1 = np.maximum(bi[0], boxes[:, 0])
369
+ yy1 = np.maximum(bi[1], boxes[:, 1])
370
+ xx2 = np.minimum(bi[2], boxes[:, 2])
371
+ yy2 = np.minimum(bi[3], boxes[:, 3])
372
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
373
+ a_i = max(1e-7, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
374
+ iou = inter / (a_i + areas - inter + 1e-7)
375
+ dup = iou > iou_thresh
376
+ dup[i] = False
377
+ # Never cross-suppress fire extinguisher vs fire/smoke.
378
+ dup &= ~((cls_ids == ext_cls) | (cls_ids[i] == ext_cls))
379
+ suppressed |= dup
380
+ keep_idx = np.array(keep, dtype=np.intp)
381
+ return boxes[keep_idx], scores[keep_idx], cls_ids[keep_idx]
382
+
383
+ @staticmethod
384
+ def _max_score_per_cluster(
385
+ post_boxes: np.ndarray,
386
+ post_cls: np.ndarray,
387
+ full_boxes: np.ndarray,
388
+ full_scores: np.ndarray,
389
+ full_cls: np.ndarray,
390
+ iou_thresh: float,
391
+ ) -> np.ndarray:
392
+ """For each kept (post-NMS) box, return the max score over the FULL
393
+ candidate set among same-class boxes with IoU >= iou_thresh.
394
+
395
+ Used after horizontal-flip TTA: a high-confidence flipped detection
396
+ can raise the score of the corresponding original detection.
397
+ """
398
+ n = len(post_boxes)
399
+ if n == 0:
400
+ return np.empty(0, dtype=np.float32)
401
+ full_areas = (np.maximum(0.0, full_boxes[:, 2] - full_boxes[:, 0]) *
402
+ np.maximum(0.0, full_boxes[:, 3] - full_boxes[:, 1]))
403
+ out = np.empty(n, dtype=np.float32)
404
+ for i in range(n):
405
+ bi = post_boxes[i]
406
+ xx1 = np.maximum(bi[0], full_boxes[:, 0])
407
+ yy1 = np.maximum(bi[1], full_boxes[:, 1])
408
+ xx2 = np.minimum(bi[2], full_boxes[:, 2])
409
+ yy2 = np.minimum(bi[3], full_boxes[:, 3])
410
+ inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
411
+ a_i = max(0.0, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
412
+ iou = inter / (a_i + full_areas - inter + 1e-7)
413
+ cluster = (iou >= iou_thresh) & (full_cls == post_cls[i])
414
+ out[i] = float(np.max(full_scores[cluster])) if np.any(cluster) else 0.0
415
+ return out
416
+
417
+ def _loose_conf_mask(
418
+ self, scores: np.ndarray, cls_ids: np.ndarray
419
+ ) -> np.ndarray:
420
+ """Pre-filter: keep candidates that could pass threshold or bonus rescue."""
421
+ floor = self._conf_thres_array[cls_ids] - self._bonus_array[cls_ids]
422
+ return scores >= floor
423
+
424
+ def _conf_filter_mask(
425
+ self, scores: np.ndarray, cls_ids: np.ndarray
426
+ ) -> np.ndarray:
427
+ """Boolean keep-mask: score >= per-class threshold, with a per-class
428
+ rescue -- admit top-1 when score >= (threshold - bonus). Runs after
429
+ sane-box filtering so tiny FPs cannot block bonus rescue."""
430
+ if len(scores) == 0:
431
+ return np.zeros(0, dtype=bool)
432
+ thr = self._conf_thres_array[cls_ids]
433
+ keep = scores >= thr
434
+ ext_cls = self._cls_fire_extinguisher
435
+ for c in np.unique(cls_ids):
436
+ b = float(self._bonus_array[c])
437
+ if b <= 0.0:
438
+ continue
439
+ cm = cls_ids == c
440
+ idx = np.where(cm)[0]
441
+ top = int(idx[int(np.argmax(scores[idx]))])
442
+ floor = float(self._conf_thres_array[c] - b)
443
+ if scores[top] < floor:
444
+ continue
445
+ if c == ext_cls:
446
+ keep[top] = True
447
+ elif not keep[cm].any():
448
+ keep[top] = True
449
+ return keep
450
+
451
+ def _filter_sane_boxes(
452
+ self,
453
+ boxes: np.ndarray,
454
+ scores: np.ndarray,
455
+ cls_ids: np.ndarray,
456
+ orig_size: tuple[int, int],
457
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
458
+ """Drop tiny / degenerate / image-spanning / extreme-AR boxes (FP)."""
459
+ if len(boxes) == 0:
460
+ return boxes, scores, cls_ids
461
+ orig_w, orig_h = orig_size
462
+ image_area = float(orig_w * orig_h)
463
+ keep = []
464
+ for i, box in enumerate(boxes):
465
+ x1, y1, x2, y2 = box.tolist()
466
+ bw = x2 - x1
467
+ bh = y2 - y1
468
+ if bw <= 0 or bh <= 0:
469
+ continue
470
+ if bw < self.min_side or bh < self.min_side:
471
+ continue
472
+ area = bw * bh
473
+ if area < self.min_box_area:
474
+ continue
475
+ if area > 0.95 * image_area:
476
+ continue
477
+ ar = max(bw / max(bh, 1e-6), bh / max(bw, 1e-6))
478
+ if ar > self.max_aspect_ratio:
479
+ continue
480
+ keep.append(i)
481
+ if not keep:
482
+ return (
483
+ np.empty((0, 4), dtype=np.float32),
484
+ np.empty((0,), dtype=np.float32),
485
+ np.empty((0,), dtype=np.int32),
486
+ )
487
+ k = np.array(keep, dtype=np.intp)
488
+ return boxes[k], scores[k], cls_ids[k]
489
+
490
+ def _per_view_pipeline(
491
+ self,
492
+ boxes: np.ndarray,
493
+ scores: np.ndarray,
494
+ cls_ids: np.ndarray,
495
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
496
+ """Per-view post-processing pipeline: per-class NMS -> nested filter -> cap -> cross-class dedup."""
497
+ if len(boxes) > 1:
498
+ keep = self._per_class_hard_nms(boxes, scores, cls_ids, self.iou_thres)
499
+ boxes, scores, cls_ids = boxes[keep], scores[keep], cls_ids[keep]
500
+ boxes, scores, cls_ids = self._nested_zone_filter(
501
+ boxes, scores, cls_ids
502
+ )
503
+ if len(scores) > self.max_det:
504
+ top = np.argsort(-scores)[: self.max_det]
505
+ boxes, scores, cls_ids = boxes[top], scores[top], cls_ids[top]
506
+ if len(boxes) > 1:
507
+ boxes, scores, cls_ids = self._cross_class_dedup_op(
508
+ boxes, scores, cls_ids, self.cross_iou_thresh
509
+ )
510
+ return boxes, scores, cls_ids
511
+
512
+ @staticmethod
513
+ def _build_results(
514
+ boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray
515
+ ) -> list[BoundingBox]:
516
+ results: list[BoundingBox] = []
517
+ for box, conf, cls_id in zip(boxes, scores, cls_ids):
518
+ x1, y1, x2, y2 = box.tolist()
519
+ if x2 <= x1 or y2 <= y1:
520
+ continue
521
+ results.append(
522
+ BoundingBox(
523
+ x1=int(math.floor(x1)),
524
+ y1=int(math.floor(y1)),
525
+ x2=int(math.ceil(x2)),
526
+ y2=int(math.ceil(y2)),
527
+ cls_id=int(cls_id),
528
+ conf=float(conf),
529
+ )
530
+ )
531
+ return results
532
+
533
+ def _decode_final_dets(
534
+ self,
535
+ preds: np.ndarray,
536
+ ratio: float,
537
+ pad: tuple[float, float],
538
+ orig_size: tuple[int, int],
539
+ ) -> list[BoundingBox]:
540
+ """Final-detection output path: rows shaped [x1, y1, x2, y2, conf, cls_id]."""
541
+ if preds.ndim == 3 and preds.shape[0] == 1:
542
+ preds = preds[0]
543
+ if preds.ndim != 2 or preds.shape[1] < 6:
544
+ raise ValueError(f"Unexpected ONNX final-det output shape: {preds.shape}")
545
+
546
+ boxes = preds[:, :4].astype(np.float32)
547
+ scores = preds[:, 4].astype(np.float32)
548
+ cls_ids = preds[:, 5].astype(np.int32)
549
+ cls_ids = self.cls_remap[cls_ids]
550
+
551
+ keep = self._loose_conf_mask(scores, cls_ids)
552
+ boxes = boxes[keep]
553
+ scores = scores[keep]
554
+ cls_ids = cls_ids[keep]
555
+ if len(boxes) == 0:
556
+ return []
557
+
558
+ pad_w, pad_h = pad
559
+ boxes[:, [0, 2]] -= pad_w
560
+ boxes[:, [1, 3]] -= pad_h
561
+ boxes /= ratio
562
+ boxes = self._clip_boxes(boxes, orig_size)
563
+
564
+ boxes, scores, cls_ids = self._filter_sane_boxes(
565
+ boxes, scores, cls_ids, orig_size
566
+ )
567
+ if len(boxes) == 0:
568
+ return []
569
+
570
+ keep = self._conf_filter_mask(scores, cls_ids)
571
+ boxes = boxes[keep]
572
+ scores = scores[keep]
573
+ cls_ids = cls_ids[keep]
574
+ if len(boxes) == 0:
575
+ return []
576
+
577
+ boxes, scores, cls_ids = self._per_view_pipeline(boxes, scores, cls_ids)
578
+ return self._build_results(boxes, scores, cls_ids)
579
+
580
+ def _decode_raw_yolo(
581
+ self,
582
+ preds: np.ndarray,
583
+ ratio: float,
584
+ pad: tuple[float, float],
585
+ orig_size: tuple[int, int],
586
+ ) -> list[BoundingBox]:
587
+ """Fallback raw-YOLO output path: per-anchor class logits."""
588
+ if preds.ndim != 3 or preds.shape[0] != 1:
589
+ raise ValueError(f"Unexpected raw ONNX output shape: {preds.shape}")
590
+ preds = preds[0]
591
+ if preds.shape[0] <= 16 and preds.shape[1] > preds.shape[0]:
592
+ preds = preds.T
593
+ if preds.ndim != 2 or preds.shape[1] < 5:
594
+ raise ValueError(f"Unexpected raw output shape: {preds.shape}")
595
+
596
+ boxes_xywh = preds[:, :4].astype(np.float32)
597
+ cls_part = preds[:, 4:].astype(np.float32)
598
+ if cls_part.shape[1] == 1:
599
+ scores = cls_part[:, 0]
600
+ cls_ids = np.zeros(len(scores), dtype=np.int32)
601
+ else:
602
+ cls_ids = np.argmax(cls_part, axis=1).astype(np.int32)
603
+ scores = cls_part[np.arange(len(cls_part)), cls_ids]
604
+ cls_ids = self.cls_remap[cls_ids]
605
+
606
+ keep = self._loose_conf_mask(scores, cls_ids)
607
+ boxes_xywh = boxes_xywh[keep]
608
+ scores = scores[keep]
609
+ cls_ids = cls_ids[keep]
610
+ if len(boxes_xywh) == 0:
611
+ return []
612
+ boxes = self._xywh_to_xyxy(boxes_xywh)
613
+
614
+ pad_w, pad_h = pad
615
+ boxes[:, [0, 2]] -= pad_w
616
+ boxes[:, [1, 3]] -= pad_h
617
+ boxes /= ratio
618
+ boxes = self._clip_boxes(boxes, orig_size)
619
+
620
+ boxes, scores, cls_ids = self._filter_sane_boxes(
621
+ boxes, scores, cls_ids, orig_size
622
+ )
623
+ if len(boxes) == 0:
624
+ return []
625
+
626
+ keep = self._conf_filter_mask(scores, cls_ids)
627
+ boxes = boxes[keep]
628
+ scores = scores[keep]
629
+ cls_ids = cls_ids[keep]
630
+ if len(boxes) == 0:
631
+ return []
632
+
633
+ boxes, scores, cls_ids = self._per_view_pipeline(boxes, scores, cls_ids)
634
+ return self._build_results(boxes, scores, cls_ids)
635
+
636
+ def _postprocess(
637
+ self,
638
+ output: np.ndarray,
639
+ ratio: float,
640
+ pad: tuple[float, float],
641
+ orig_size: tuple[int, int],
642
+ ) -> list[BoundingBox]:
643
+ if output.ndim == 2 and output.shape[1] >= 6:
644
+ return self._decode_final_dets(output, ratio, pad, orig_size)
645
+ if output.ndim == 3 and output.shape[0] == 1 and output.shape[2] == 6:
646
+ return self._decode_final_dets(output, ratio, pad, orig_size)
647
+ return self._decode_raw_yolo(output, ratio, pad, orig_size)
648
+
649
+ def _predict_single(self, image: np.ndarray) -> list[BoundingBox]:
650
+ if image is None:
651
+ raise ValueError("Input image is None")
652
+ if not isinstance(image, np.ndarray):
653
+ raise TypeError(f"Input is not numpy array: {type(image)}")
654
+ if image.ndim != 3:
655
+ raise ValueError(f"Expected HWC image, got shape={image.shape}")
656
+ if image.shape[0] <= 0 or image.shape[1] <= 0:
657
+ raise ValueError(f"Invalid image shape={image.shape}")
658
+ if image.shape[2] != 3:
659
+ raise ValueError(f"Expected 3 channels, got shape={image.shape}")
660
+ if image.dtype != np.uint8:
661
+ image = image.astype(np.uint8)
662
+
663
+ input_tensor, ratio, pad, orig_size = self._preprocess(image)
664
+ expected = (1, 3, self.input_height, self.input_width)
665
+ if input_tensor.shape != expected:
666
+ raise ValueError(
667
+ f"Bad input tensor shape={input_tensor.shape}, expected={expected}"
668
+ )
669
+
670
+ outputs = self.session.run(self.output_names, {self.input_name: input_tensor})
671
+ return self._postprocess(outputs[0], ratio, pad, orig_size)
672
+
673
+ def _predict_tta(self, image: np.ndarray) -> list[BoundingBox]:
674
+ """Horizontal-flip TTA.
675
+
676
+ Strategy:
677
+ 1. Predict on original and on flipped image.
678
+ 2. Map flipped boxes back to original coordinates.
679
+ 3. Per-class hard NMS on the union.
680
+ 4. For each kept box, compute the max same-class score across the
681
+ FULL union (not just the post-NMS subset) -- this lets a high-
682
+ confidence flipped detection raise a borderline original one.
683
+ 5. Cross-class dedup to suppress same-physical-object multi-class.
684
+ """
685
+ boxes_orig = self._predict_single(image)
686
+ flipped = cv2.flip(image, 1)
687
+ boxes_flip = self._predict_single(flipped)
688
+ w = image.shape[1]
689
+ boxes_flip = [
690
+ BoundingBox(
691
+ x1=w - b.x2, y1=b.y1, x2=w - b.x1, y2=b.y2,
692
+ cls_id=b.cls_id, conf=b.conf,
693
+ )
694
+ for b in boxes_flip
695
+ ]
696
+ all_boxes = boxes_orig + boxes_flip
697
+ if not all_boxes:
698
+ return []
699
+
700
+ coords = np.array(
701
+ [[b.x1, b.y1, b.x2, b.y2] for b in all_boxes], dtype=np.float32
702
+ )
703
+ scores = np.array([b.conf for b in all_boxes], dtype=np.float32)
704
+ cls_ids = np.array([b.cls_id for b in all_boxes], dtype=np.int32)
705
+
706
+ hard_keep = self._per_class_hard_nms(coords, scores, cls_ids, self.iou_thres)
707
+ if len(hard_keep) == 0:
708
+ return []
709
+ kept_coords = coords[hard_keep]
710
+ kept_scores = scores[hard_keep]
711
+ kept_cls = cls_ids[hard_keep]
712
+ kept_coords, kept_scores, kept_cls = self._nested_zone_filter(
713
+ kept_coords, kept_scores, kept_cls
714
+ )
715
+ if len(kept_coords) == 0:
716
+ return []
717
+ if len(kept_scores) > self.max_det:
718
+ top = np.argsort(-kept_scores)[: self.max_det]
719
+ kept_coords = kept_coords[top]
720
+ kept_scores = kept_scores[top]
721
+ kept_cls = kept_cls[top]
722
+
723
+ boosted = self._max_score_per_cluster(
724
+ kept_coords, kept_cls,
725
+ coords, scores, cls_ids, self.iou_thres,
726
+ )
727
+ if len(kept_coords) > 1:
728
+ kept_coords, boosted, kept_cls = self._cross_class_dedup_op(
729
+ kept_coords, boosted, kept_cls, self.cross_iou_thresh
730
+ )
731
+
732
+ return [
733
+ BoundingBox(
734
+ x1=int(math.floor(kept_coords[j, 0])),
735
+ y1=int(math.floor(kept_coords[j, 1])),
736
+ x2=int(math.ceil(kept_coords[j, 2])),
737
+ y2=int(math.ceil(kept_coords[j, 3])),
738
+ cls_id=int(kept_cls[j]),
739
+ conf=float(boosted[j]),
740
+ )
741
+ for j in range(len(kept_coords))
742
+ ]
743
+
744
+ def predict_batch(
745
+ self,
746
+ batch_images: list[ndarray],
747
+ offset: int,
748
+ n_keypoints: int,
749
+ ) -> list[TVFrameResult]:
750
+ results: list[TVFrameResult] = []
751
+ for frame_number_in_batch, image in enumerate(batch_images):
752
+ try:
753
+ if self.use_tta:
754
+ boxes = self._predict_tta(image)
755
+ else:
756
+ boxes = self._predict_single(image)
757
+ except Exception as e:
758
+ print(
759
+ f"⚠️ Inference failed for frame "
760
+ f"{offset + frame_number_in_batch}: {e}"
761
+ )
762
+ boxes = []
763
+ results.append(
764
+ TVFrameResult(
765
+ frame_id=offset + frame_number_in_batch,
766
+ boxes=boxes,
767
+ keypoints=[(0, 0) for _ in range(max(0, int(n_keypoints)))],
768
+ )
769
+ )
770
+ return results
 
 
 
 
 
 
 
weights.onnx CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:5e890080398746b8365fa969ec087690767b0eae1a9b528d0e7bd6d0fff2c1ed
3
- size 19407372
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:6e64f8b16bbaa811aa577f962ba0a961698fef8b6f6d72e788df58687359c060
3
+ size 19407374