from pathlib import Path import math import cv2 import numpy as np import onnxruntime as ort from numpy import ndarray from pydantic import BaseModel class BoundingBox(BaseModel): x1: int y1: int x2: int y2: int cls_id: int conf: float class TVFrameResult(BaseModel): frame_id: int boxes: list[BoundingBox] keypoints: list[tuple[int, int]] class Miner: class_names = ["fire", "smoke", "fire extinguisher"] _model_class_order = ["fire", "fire extinguisher", "smoke"] _parent_split_class_names = ("fire", "fire extinguisher") conf_thresh = 0.000005 isolated_conf_thresh = 0.008 same_class_iou_thresh = 0.6 edge_conf_thresh = 0.005 edge_tol = 1.0 overlap_low_conf_thresh = 0.003 overlap_same_class_iou_thresh = 0.5 group_iou_thresh = 0.7 flip_low_conf_thresh = 0.2 fire_ext_same_object_overlap_thresh = 0.7 fire_ext_group_involve_thresh = 0.9 smoke_involve_thresh = 0.7 smoke_prefer_larger_conf_thresh = 0.3 smoke_cluster_merge_involve_thresh = 0.3 smoke_cluster_merge_min_size = 2 smoke_cluster_merge_early_min_count = 3 small_group_min_count = 2 small_group_low_conf_thresh = 0.02 fire_parent_min_small_count = 2 fire_parent_involve_thresh = 0.7 fire_split_inner_involve_thresh = 0.5 fire_split_max_inner_area_ratio = 0.5 fire_split_cluster_iou_thresh = 0.5 fire_split_min_clusters = 2 fire_split_min_inner_conf = 0.08 fire_split_peak_member_min_conf = 0.05 fire_split_peak_dist_ratio = 0.10 fire_split_peak_dist_min_px = 12.0 fire_split_peak_dist_max_px = 48.0 fire_split_centroid_preserve_px = 12.0 fire_split_conf_boost = 0.2 fire_split_parent_blob_conf_thresh = 0.5 fire_min_conf = 0.165 smoke_min_conf = 0.07 fire_ext_min_conf = 0.037 fire_ext_color_filter_max_conf = 0.20 fire_ext_min_red_dom_frac = 0.03 fire_ext_color_min_mean_r = 50.0 fire_ext_min_r_minus_g = 0.0 smoke_anchor_min_smoke_conf = 0.63 smoke_anchor_base_height_frac = 0.35 smoke_anchor_crop_width_pad_ratio = 0.15 smoke_anchor_crop_down_pad_px = 24.0 smoke_anchor_corroborate_involve_thresh = 0.5 smoke_anchor_fire_floor_conf = 0.165 smoke_anchor_min_fire_width_px = 10.0 smoke_anchor_min_fire_height_px = 10.0 smoke_anchor_strong_smoke_conf = 0.7 smoke_anchor_corroborated_min_fire_width_px = 4.0 smoke_anchor_corroborated_min_fire_height_px = 4.0 smoke_anchor_max_fire_area_ratio = 0.25 smoke_anchor_corroborate_area_ratio_max = 4.0 smoke_anchor_max_fires_per_smoke = 2 smoke_anchor_max_crop_fires = 2 smoke_anchor_max_probe_smokes = 2 smoke_anchor_separate_fire_involve_thresh = 0.5 smoke_anchor_max_probe_fire_width_px = 52.0 smoke_anchor_max_probe_fire_height_px = 48.0 smoke_anchor_max_probe_fire_area_ratio = 0.04 smoke_anchor_max_probe_crop_width_ratio = 0.25 smoke_anchor_max_probe_crop_height_ratio = 0.60 smoke_anchor_max_probe_fire_bottom_dist_frac = 0.10 smoke_anchor_probe_rescue_min_conf = 0.000005 smoke_anchor_corroborated_probe_min_conf = 0.00005 smoke_anchor_corroborated_refine_pad_px = 3.0 smoke_anchor_expanded_width_ratio = 1.25 ext_probe_min_cluster_count = 2 ext_probe_cluster_centroid_dist_px = 28.0 ext_probe_min_cluster_best_conf = 0.003 ext_probe_crop_pad_ratio = 0.85 ext_probe_min_crop_size_px = 200.0 ext_probe_crop_rescue_min_conf = 0.5 ext_probe_crop_rescue_min_width_px = 6.0 ext_probe_rescue_min_conf = 0.000005 ext_probe_min_width_px = 10.0 ext_probe_min_height_px = 10.0 ext_probe_max_width_px = 120.0 ext_probe_max_height_px = 160.0 ext_probe_fire_exclude_involve_thresh = 0.5 ext_probe_near_fire_gap_px = 24.0 ext_probe_duplicate_involve_thresh = 0.5 ext_probe_max_per_frame = 2 ext_probe_edge_tol = 1.0 ext_probe_singleton_min_conf = 0.00005 ext_probe_singleton_max_area = 500.0 ext_probe_singleton_max_width_px = 20.0 ext_probe_singleton_max_height_px = 35.0 ext_probe_singleton_min_edge_margin_px = 12.0 smoke_anchor_probe_min_mean_r = 65.0 smoke_anchor_probe_min_max_rgb = 80.0 smoke_anchor_probe_min_warm_frac = 0.05 smoke_anchor_probe_warm_min_max_rgb = 120.0 smoke_anchor_probe_warm_min_mean_r = 120.0 smoke_anchor_probe_strong_warm_frac = 0.15 smoke_anchor_probe_min_bright_frac = 0.12 smoke_anchor_probe_bright_thresh = 150 smoke_anchor_probe_min_hot_max_rgb = 180.0 smoke_anchor_probe_min_r_minus_g = 2.0 smoke_anchor_probe_sparse_hot_min_max_rgb = 200.0 smoke_anchor_probe_sparse_hot_min_bright_frac = 0.01 smoke_expand_min_seed_conf = 0.5 smoke_expand_crop_width_pad_ratio = 0.45 smoke_expand_crop_up_pad_ratio = 0.15 smoke_expand_crop_down_pad_px = 32.0 smoke_expand_corroborate_involve_thresh = 0.3 smoke_expand_vertical_band_pad_ratio = 0.10 smoke_expand_band_max_horizontal_gap_ratio = 0.25 smoke_expand_upward_diagonal_max_gap_ratio = 0.45 smoke_expand_min_probe_smoke_conf = 0.10 smoke_expand_min_tta_smoke_conf = 0.12 smoke_expand_min_tta_extend_conf = 0.15 smoke_expand_tta_view_iou_thresh = 0.5 smoke_expand_wide_seed_max_width_ratio = 0.25 smoke_expand_upper_plume_max_y_ratio = 0.72 smoke_expand_wide_crop_x_start_ratio = 0.05 smoke_expand_wide_crop_x_end_ratio = 0.95 smoke_expand_min_crop_width_ratio = 0.60 smoke_expand_border_strip_ratio = 0.15 smoke_expand_max_color_dist = 45.0 smoke_expand_max_frame_ratio = 0.20 smoke_expand_min_extend_probe_conf = 0.15 smoke_expand_min_bidi_crop_conf = 0.40 smoke_expand_min_upward_tta_conf = 0.012 smoke_expand_min_upward_tta_involve_conf = 0.0002 smoke_expand_min_upward_crop_conf = 0.012 smoke_expand_max_seeds = 1 smoke_expand_skip_min_conf = 0.85 smoke_expand_skip_min_width_ratio = 0.15 smoke_expand_skip_max_y1_ratio = 0.04 def __init__(self, path_hf_repo: Path) -> None: model_path = path_hf_repo / "weights.onnx" self.group_object_counts: list[int] = [] self.group_max_confidences: list[float] = [] self.group_avg_confidences: list[float] = [] self.cls_remap = np.array( [self.class_names.index(n) for n in self._model_class_order], dtype=np.int32, ) print("ORT version:", ort.__version__) try: ort.preload_dlls() print("✅ onnxruntime.preload_dlls() success") except Exception as e: print(f"⚠️ preload_dlls failed: {e}") print("ORT available providers BEFORE session:", ort.get_available_providers()) sess_options = ort.SessionOptions() sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL try: self.session = ort.InferenceSession( str(model_path), sess_options=sess_options, providers=["CUDAExecutionProvider", "CPUExecutionProvider"], ) print("✅ Created ORT session with preferred CUDA provider list") except Exception as e: print(f"⚠️ CUDA session creation failed, falling back to CPU: {e}") self.session = ort.InferenceSession( str(model_path), sess_options=sess_options, providers=["CPUExecutionProvider"], ) print("ORT session providers:", self.session.get_providers()) for inp in self.session.get_inputs(): print("INPUT:", inp.name, inp.shape, inp.type) for out in self.session.get_outputs(): print("OUTPUT:", out.name, out.shape, out.type) self.input_name = self.session.get_inputs()[0].name self.output_names = [output.name for output in self.session.get_outputs()] self.input_shape = self.session.get_inputs()[0].shape self.input_height = self._safe_dim(self.input_shape[2], default=1280) self.input_width = self._safe_dim(self.input_shape[3], default=1280) self.use_tta = True print(f"✅ ONNX model loaded from: {model_path}") print(f"✅ ONNX providers: {self.session.get_providers()}") print(f"✅ ONNX input: name={self.input_name}, shape={self.input_shape}") def __repr__(self) -> str: return ( f"ONNXRuntime(session={type(self.session).__name__}, " f"providers={self.session.get_providers()})" ) @staticmethod def _safe_dim(value, default: int) -> int: return value if isinstance(value, int) and value > 0 else default def _letterbox( self, image: ndarray, new_shape: tuple[int, int], color=(114, 114, 114), ) -> tuple[ndarray, float, tuple[float, float]]: h, w = image.shape[:2] new_w, new_h = new_shape ratio = min(new_w / w, new_h / h) resized_w = int(round(w * ratio)) resized_h = int(round(h * ratio)) if (resized_w, resized_h) != (w, h): interp = cv2.INTER_CUBIC if ratio > 1.0 else cv2.INTER_LINEAR image = cv2.resize(image, (resized_w, resized_h), interpolation=interp) dw = (new_w - resized_w) / 2.0 dh = (new_h - resized_h) / 2.0 left = int(round(dw - 0.1)) right = int(round(dw + 0.1)) top = int(round(dh - 0.1)) bottom = int(round(dh + 0.1)) padded = cv2.copyMakeBorder( image, top, bottom, left, right, borderType=cv2.BORDER_CONSTANT, value=color, ) return padded, ratio, (dw, dh) def _preprocess( self, image: ndarray ) -> tuple[np.ndarray, float, tuple[float, float], tuple[int, int]]: orig_h, orig_w = image.shape[:2] img, ratio, pad = self._letterbox( image, (self.input_width, self.input_height) ) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = img.astype(np.float32) / 255.0 img = np.transpose(img, (2, 0, 1))[None, ...] img = np.ascontiguousarray(img, dtype=np.float32) return img, ratio, pad, (orig_w, orig_h) @staticmethod def _clip_boxes(boxes: np.ndarray, image_size: tuple[int, int]) -> np.ndarray: w, h = image_size boxes[:, 0] = np.clip(boxes[:, 0], 0, w - 1) boxes[:, 1] = np.clip(boxes[:, 1], 0, h - 1) boxes[:, 2] = np.clip(boxes[:, 2], 0, w - 1) boxes[:, 3] = np.clip(boxes[:, 3], 0, h - 1) return boxes @staticmethod def _build_results( boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray ) -> list[BoundingBox]: results: list[BoundingBox] = [] for box, conf, cls_id in zip(boxes, scores, cls_ids): x1, y1, x2, y2 = box.tolist() if x2 <= x1 or y2 <= y1: continue 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), ) ) return results def _map_boxes_to_orig( self, boxes: np.ndarray, ratio: float, pad: tuple[float, float], orig_size: tuple[int, int], ) -> np.ndarray: pad_w, pad_h = pad boxes = boxes.copy() boxes[:, [0, 2]] -= pad_w boxes[:, [1, 3]] -= pad_h boxes /= ratio return self._clip_boxes(boxes, orig_size) @staticmethod def _compute_iou_matrix(boxes: np.ndarray) -> np.ndarray: n = len(boxes) if n == 0: return np.empty((0, 0), dtype=np.float32) areas = (np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) * np.maximum(0.0, boxes[:, 3] - boxes[:, 1])) iou = np.zeros((n, n), dtype=np.float32) for i in range(n): bi = boxes[i] xx1 = np.maximum(bi[0], boxes[i + 1:, 0]) yy1 = np.maximum(bi[1], boxes[i + 1:, 1]) xx2 = np.minimum(bi[2], boxes[i + 1:, 2]) yy2 = np.minimum(bi[3], boxes[i + 1:, 3]) inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1) a_i = max(1e-7, float((bi[2] - bi[0]) * (bi[3] - bi[1]))) iou[i, i + 1:] = inter / (a_i + areas[i + 1:] - inter + 1e-7) iou = iou + iou.T np.fill_diagonal(iou, 1.0) return iou def _group_overlapping_detections( self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray, iou_thresh: float | None = None, view_ids: np.ndarray | None = None, record_stats: bool = True, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: if iou_thresh is None: iou_thresh = self.group_iou_thresh n = len(boxes) if record_stats: self.group_object_counts = [] self.group_max_confidences = [] self.group_avg_confidences = [] if n == 0: return boxes, scores, cls_ids boxes = np.asarray(boxes, dtype=np.float32) scores = np.asarray(scores, dtype=np.float32) cls_ids = np.asarray(cls_ids, dtype=np.int32) if view_ids is not None: view_ids = np.asarray(view_ids, dtype=np.int32) parent = list(range(n)) def find(x: int) -> int: while parent[x] != x: parent[x] = parent[parent[x]] x = parent[x] return x def union(a: int, b: int) -> None: ra, rb = find(a), find(b) if ra != rb: parent[rb] = ra iou = self._compute_iou_matrix(boxes) for i in range(n): for j in range(i + 1, n): if cls_ids[i] != cls_ids[j]: continue if view_ids is not None and view_ids[i] != view_ids[j]: continue if iou[i, j] > iou_thresh: union(i, j) groups: dict[int, list[int]] = {} for i in range(n): groups.setdefault(find(i), []).append(i) kept_groups: list[tuple[int, int, float, float]] = [] for members in groups.values(): member_scores = scores[members] count = len(members) max_conf = float(np.max(member_scores)) avg_conf = float(np.mean(member_scores)) best = members[int(np.argmax(member_scores))] kept_groups.append((best, count, max_conf, avg_conf)) kept_groups.sort(key=lambda item: item[0]) keep = np.array([item[0] for item in kept_groups], dtype=np.intp) if record_stats: self.group_object_counts = [item[1] for item in kept_groups] self.group_max_confidences = [item[2] for item in kept_groups] self.group_avg_confidences = [item[3] for item in kept_groups] return boxes[keep], scores[keep], cls_ids[keep] def _remove_isolated_low_conf_same_class( self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray, view_ids: np.ndarray | None = None, ) -> tuple[np.ndarray, np.ndarray, np.ndarray] | tuple[ np.ndarray, np.ndarray, np.ndarray, np.ndarray ]: n = len(boxes) if n == 0: if view_ids is None: return boxes, scores, cls_ids return boxes, scores, cls_ids, view_ids if n == 1: if scores[0] < self.isolated_conf_thresh: empty_boxes = np.empty((0, 4), dtype=np.float32) empty_scores = np.empty((0,), dtype=np.float32) empty_cls = np.empty((0,), dtype=np.int32) if view_ids is None: return empty_boxes, empty_scores, empty_cls return empty_boxes, empty_scores, empty_cls, np.empty((0,), dtype=np.int32) if view_ids is None: return boxes, scores, cls_ids return boxes, scores, cls_ids, view_ids boxes = np.asarray(boxes, dtype=np.float32) scores = np.asarray(scores, dtype=np.float32) cls_ids = np.asarray(cls_ids, dtype=np.int32) areas = (np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) * np.maximum(0.0, boxes[:, 3] - boxes[:, 1])) keep = np.ones(n, dtype=bool) low_conf = scores < self.isolated_conf_thresh if not np.any(low_conf): if view_ids is None: return boxes, scores, cls_ids return boxes, scores, cls_ids, view_ids for i in np.where(low_conf)[0]: bi = boxes[i] xx1 = np.maximum(bi[0], boxes[:, 0]) yy1 = np.maximum(bi[1], boxes[:, 1]) xx2 = np.minimum(bi[2], boxes[:, 2]) yy2 = np.minimum(bi[3], boxes[:, 3]) inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1) a_i = max(1e-7, float((bi[2] - bi[0]) * (bi[3] - bi[1]))) iou = inter / (a_i + areas - inter + 1e-7) same_class = cls_ids == cls_ids[i] other = np.arange(n) != i if not np.any((iou > self.same_class_iou_thresh) & same_class & other): keep[i] = False if view_ids is None: return boxes[keep], scores[keep], cls_ids[keep] return boxes[keep], scores[keep], cls_ids[keep], view_ids[keep] def _remove_edge_low_conf( self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray, orig_size: tuple[int, int], view_ids: np.ndarray | None = None, ) -> tuple[np.ndarray, np.ndarray, np.ndarray] | tuple[ np.ndarray, np.ndarray, np.ndarray, np.ndarray ]: if len(boxes) == 0: if view_ids is None: return boxes, scores, cls_ids return boxes, scores, cls_ids, view_ids orig_w, orig_h = orig_size tol = self.edge_tol on_edge = ( (boxes[:, 0] <= tol) | (boxes[:, 1] <= tol) | (boxes[:, 2] >= orig_w - 1 - tol) | (boxes[:, 3] >= orig_h - 1 - tol) ) keep = ~(on_edge & (scores < self.edge_conf_thresh)) if view_ids is None: return boxes[keep], scores[keep], cls_ids[keep] return boxes[keep], scores[keep], cls_ids[keep], view_ids[keep] def _remove_overlapping_low_conf_same_class( self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray, view_ids: np.ndarray | None = None, ) -> tuple[np.ndarray, np.ndarray, np.ndarray] | tuple[ np.ndarray, np.ndarray, np.ndarray, np.ndarray ]: n = len(boxes) if n <= 1: if view_ids is None: return boxes, scores, cls_ids return boxes, scores, cls_ids, view_ids boxes = np.asarray(boxes, dtype=np.float32) scores = np.asarray(scores, dtype=np.float32) cls_ids = np.asarray(cls_ids, dtype=np.int32) areas = (np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) * np.maximum(0.0, boxes[:, 3] - boxes[:, 1])) keep = np.ones(n, dtype=bool) low_conf = scores < self.overlap_low_conf_thresh if not np.any(low_conf): if view_ids is None: return boxes, scores, cls_ids return boxes, scores, cls_ids, view_ids for i in np.where(low_conf)[0]: bi = boxes[i] xx1 = np.maximum(bi[0], boxes[:, 0]) yy1 = np.maximum(bi[1], boxes[:, 1]) xx2 = np.minimum(bi[2], boxes[:, 2]) yy2 = np.minimum(bi[3], boxes[:, 3]) inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1) a_i = max(1e-7, float((bi[2] - bi[0]) * (bi[3] - bi[1]))) iou = inter / (a_i + areas - inter + 1e-7) same_class = cls_ids == cls_ids[i] other = np.arange(n) != i if np.any((iou >= self.overlap_same_class_iou_thresh) & same_class & other): keep[i] = False if view_ids is None: return boxes[keep], scores[keep], cls_ids[keep] return boxes[keep], scores[keep], cls_ids[keep], view_ids[keep] def _decode_preds_to_arrays( self, preds: np.ndarray, ratio: float, pad: tuple[float, float], orig_size: tuple[int, int], ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: if preds.ndim == 3 and preds.shape[0] == 1: preds = preds[0] if preds.ndim != 2 or preds.shape[1] < 6: raise ValueError(f"Unexpected ONNX final-det output shape: {preds.shape}") boxes = preds[:, :4].astype(np.float32) scores = preds[:, 4].astype(np.float32) cls_ids = self.cls_remap[preds[:, 5].astype(np.int32)] keep = scores >= self.conf_thresh boxes = boxes[keep] scores = scores[keep] cls_ids = cls_ids[keep] if len(boxes) == 0: return ( np.empty((0, 4), dtype=np.float32), np.empty((0,), dtype=np.float32), np.empty((0,), dtype=np.int32), ) boxes = self._map_boxes_to_orig(boxes, ratio, pad, orig_size) valid = (boxes[:, 2] > boxes[:, 0]) & (boxes[:, 3] > boxes[:, 1]) return boxes[valid], scores[valid], cls_ids[valid] @staticmethod def _map_flipped_boxes_arrays( boxes: np.ndarray, orig_size: tuple[int, int], flip_code: int, ) -> np.ndarray: orig_w, _orig_h = orig_size out = boxes.copy() if flip_code == 1: x1 = orig_w - out[:, 2] x2 = orig_w - out[:, 0] out[:, 0] = x1 out[:, 2] = x2 else: raise ValueError(f"Unsupported flip_code: {flip_code}") return out def _remove_unconfirmed_single_view_low_conf( self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray, view_ids: np.ndarray, conf_thresh: float | None = None, match_iou_thresh: float | None = None, protect_mask: np.ndarray | None = None, ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: if conf_thresh is None: conf_thresh = self.flip_low_conf_thresh if match_iou_thresh is None: match_iou_thresh = self.overlap_same_class_iou_thresh n = len(boxes) if n == 0: return boxes, scores, cls_ids, view_ids low_conf = scores < conf_thresh if not np.any(low_conf): return boxes, scores, cls_ids, view_ids iou = self._compute_iou_matrix(boxes) keep = np.ones(n, dtype=bool) for i in np.where(low_conf)[0]: if protect_mask is not None and protect_mask[i]: continue other = np.arange(n) != i other_view = view_ids != view_ids[i] same_class = cls_ids == cls_ids[i] corroborated = np.any( other & other_view & same_class & (iou[i] >= match_iou_thresh) ) if not corroborated: keep[i] = False return boxes[keep], scores[keep], cls_ids[keep], view_ids[keep] @staticmethod def _boxes_involve_each_other( boxes: np.ndarray, i: int, j: int, overlap_thresh: float ) -> bool: bi, bj = boxes[i], boxes[j] xx1 = max(float(bi[0]), float(bj[0])) yy1 = max(float(bi[1]), float(bj[1])) xx2 = min(float(bi[2]), float(bj[2])) yy2 = min(float(bi[3]), float(bj[3])) inter = max(0.0, xx2 - xx1) * max(0.0, yy2 - yy1) if inter <= 0.0: return False area_i = max(1e-7, float((bi[2] - bi[0]) * (bi[3] - bi[1]))) area_j = max(1e-7, float((bj[2] - bj[0]) * (bj[3] - bj[1]))) return (inter / min(area_i, area_j)) >= overlap_thresh @staticmethod def _box_area(box: np.ndarray) -> float: return max(1e-7, float((box[2] - box[0]) * (box[3] - box[1]))) @staticmethod def _box_intersection_area(box_a: np.ndarray, box_b: np.ndarray) -> float: xx1 = max(float(box_a[0]), float(box_b[0])) yy1 = max(float(box_a[1]), float(box_b[1])) xx2 = min(float(box_a[2]), float(box_b[2])) yy2 = min(float(box_a[3]), float(box_b[3])) return max(0.0, xx2 - xx1) * max(0.0, yy2 - yy1) @staticmethod def _larger_involves_smaller( boxes: np.ndarray, big_idx: int, small_idx: int, involve_thresh: float ) -> bool: inter = Miner._box_intersection_area(boxes[big_idx], boxes[small_idx]) if inter <= 0.0: return False area_small = Miner._box_area(boxes[small_idx]) return (inter / area_small) >= involve_thresh @staticmethod def _either_box_involves_other( boxes: np.ndarray, i: int, j: int, involve_thresh: float ) -> bool: inter = Miner._box_intersection_area(boxes[i], boxes[j]) if inter <= 0.0: return False area_i = Miner._box_area(boxes[i]) area_j = Miner._box_area(boxes[j]) return (inter / area_i >= involve_thresh) or (inter / area_j >= involve_thresh) @staticmethod def _boxes_involve_each_other_arrays( box_a: np.ndarray, box_b: np.ndarray, involve_thresh: float ) -> bool: inter = Miner._box_intersection_area(box_a, box_b) if inter <= 0.0: return False area_a = Miner._box_area(box_a) area_b = Miner._box_area(box_b) return (inter / area_a >= involve_thresh) or (inter / area_b >= involve_thresh) def _resolve_smoke_overlap( self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray, involve_thresh: float | None = None, prefer_larger_conf_thresh: float | None = None, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: if involve_thresh is None: involve_thresh = self.smoke_involve_thresh if prefer_larger_conf_thresh is None: prefer_larger_conf_thresh = self.smoke_prefer_larger_conf_thresh n = len(boxes) if n <= 1: return boxes, scores, cls_ids cls_smoke = self.class_names.index("smoke") keep = np.ones(n, dtype=bool) for i in range(n): if not keep[i] or int(cls_ids[i]) != cls_smoke: continue for j in range(i + 1, n): if not keep[j] or int(cls_ids[j]) != cls_smoke: continue area_i = self._box_area(boxes[i]) area_j = self._box_area(boxes[j]) if area_i == area_j: big_idx, small_idx = (i, j) if scores[i] >= scores[j] else (j, i) elif area_i > area_j: big_idx, small_idx = i, j else: big_idx, small_idx = j, i if not self._larger_involves_smaller( boxes, big_idx, small_idx, involve_thresh ): continue big_conf = float(scores[big_idx]) small_conf = float(scores[small_idx]) if ( big_conf > prefer_larger_conf_thresh and small_conf > prefer_larger_conf_thresh ): keep[small_idx] = False elif big_conf >= small_conf: keep[small_idx] = False else: keep[big_idx] = False if big_idx == i: break return boxes[keep], scores[keep], cls_ids[keep] def _merge_overlapping_smoke_clusters( self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray, involve_thresh: float | None = None, min_cluster_size: int | None = None, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: if involve_thresh is None: involve_thresh = self.smoke_cluster_merge_involve_thresh if min_cluster_size is None: min_cluster_size = self.smoke_cluster_merge_min_size n = len(boxes) if n == 0: return boxes, scores, cls_ids cls_smoke = self.class_names.index("smoke") smoke_indices = [ i for i in range(n) if int(cls_ids[i]) == cls_smoke ] if len(smoke_indices) < min_cluster_size: return boxes, scores, cls_ids n_smoke = len(smoke_indices) parent = list(range(n_smoke)) def find(x: int) -> int: while parent[x] != x: parent[x] = parent[parent[x]] x = parent[x] return x def union(a: int, b: int) -> None: ra, rb = find(a), find(b) if ra != rb: parent[rb] = ra for ai in range(n_smoke): for bi in range(ai + 1, n_smoke): idx_i = smoke_indices[ai] idx_j = smoke_indices[bi] if self._either_box_involves_other( boxes, idx_i, idx_j, involve_thresh ): union(ai, bi) clusters: dict[int, list[int]] = {} for ai in range(n_smoke): root = find(ai) clusters.setdefault(root, []).append(smoke_indices[ai]) remove_indices: set[int] = set() merged_boxes: list[np.ndarray] = [] merged_scores: list[float] = [] for cluster in clusters.values(): if len(cluster) < min_cluster_size: continue remove_indices.update(cluster) union_box = boxes[cluster[0]].copy() max_conf = float(scores[cluster[0]]) for idx in cluster[1:]: union_box[0] = min(union_box[0], boxes[idx][0]) union_box[1] = min(union_box[1], boxes[idx][1]) union_box[2] = max(union_box[2], boxes[idx][2]) union_box[3] = max(union_box[3], boxes[idx][3]) max_conf = max(max_conf, float(scores[idx])) merged_boxes.append(union_box) merged_scores.append(max_conf) if not remove_indices: return boxes, scores, cls_ids keep = np.ones(n, dtype=bool) for idx in remove_indices: keep[idx] = False out_boxes = list(boxes[keep]) out_scores = [float(s) for s in scores[keep]] out_cls = [int(c) for c in cls_ids[keep]] out_boxes.extend(merged_boxes) out_scores.extend(merged_scores) out_cls.extend([cls_smoke] * len(merged_boxes)) return ( np.asarray(out_boxes, dtype=np.float32), np.asarray(out_scores, dtype=np.float32), np.asarray(out_cls, dtype=np.int32), ) def _remove_class_min_conf( self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray, fire_min_conf: float | None = None, smoke_min_conf: float | None = None, fire_ext_min_conf: float | None = None, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: if fire_min_conf is None: fire_min_conf = self.fire_min_conf if smoke_min_conf is None: smoke_min_conf = self.smoke_min_conf if fire_ext_min_conf is None: fire_ext_min_conf = self.fire_ext_min_conf n = len(boxes) if n == 0: return boxes, scores, cls_ids cls_fire = self.class_names.index("fire") cls_smoke = self.class_names.index("smoke") cls_fire_ext = self.class_names.index("fire extinguisher") keep = np.ones(n, dtype=bool) for i in range(n): cls_id = int(cls_ids[i]) if cls_id == cls_fire and float(scores[i]) < fire_min_conf: keep[i] = False elif cls_id == cls_smoke and float(scores[i]) < smoke_min_conf: keep[i] = False elif cls_id == cls_fire_ext and float(scores[i]) < fire_ext_min_conf: keep[i] = False return boxes[keep], scores[keep], cls_ids[keep] def _filter_results_min_conf( self, results: list[BoundingBox] ) -> list[BoundingBox]: if not results: return results cls_fire = self.class_names.index("fire") cls_smoke = self.class_names.index("smoke") cls_fire_ext = self.class_names.index("fire extinguisher") filtered: list[BoundingBox] = [] for box in results: if box.cls_id == cls_fire and box.conf < self.fire_min_conf: continue if box.cls_id == cls_smoke and box.conf < self.smoke_min_conf: continue if box.cls_id == cls_fire_ext and box.conf < self.fire_ext_min_conf: continue filtered.append(box) return filtered def _resolve_fire_and_ext_same_class_overlap( self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray, overlap_thresh: float | None = None, group_involve_thresh: float | None = None, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: if overlap_thresh is None: overlap_thresh = self.fire_ext_same_object_overlap_thresh if group_involve_thresh is None: group_involve_thresh = self.fire_ext_group_involve_thresh n = len(boxes) if n <= 1: return boxes, scores, cls_ids cls_fire = self.class_names.index("fire") dedupe_classes = { cls_fire, self.class_names.index("fire extinguisher"), } keep = np.ones(n, dtype=bool) for i in range(n): if not keep[i] or int(cls_ids[i]) not in dedupe_classes: continue for j in range(i + 1, n): if not keep[j] or cls_ids[i] != cls_ids[j]: continue should_dedupe = self._either_box_involves_other( boxes, i, j, group_involve_thresh ) if not should_dedupe: if not self._boxes_involve_each_other( boxes, i, j, overlap_thresh ): continue if int(cls_ids[i]) == cls_fire: cx_i, cy_i = self._box_centroid(boxes[i]) cx_j, cy_j = self._box_centroid(boxes[j]) dx = cx_i - cx_j dy = cy_i - cy_j if ( dx * dx + dy * dy > self.fire_split_centroid_preserve_px * self.fire_split_centroid_preserve_px ): continue should_dedupe = True if not should_dedupe: continue if scores[i] >= scores[j]: keep[j] = False else: keep[i] = False break return boxes[keep], scores[keep], cls_ids[keep] def _parent_split_cls_ids(self) -> list[int]: return [self.class_names.index(name) for name in self._parent_split_class_names] def _remove_spanning_parent_fire( self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray, view_ids: np.ndarray | None = None, involve_thresh: float | None = None, min_small_count: int | None = None, ) -> tuple[np.ndarray, np.ndarray, np.ndarray] | tuple[ np.ndarray, np.ndarray, np.ndarray, np.ndarray ]: if involve_thresh is None: involve_thresh = self.fire_parent_involve_thresh if min_small_count is None: min_small_count = self.fire_parent_min_small_count n = len(boxes) if n <= 1: if view_ids is None: return boxes, scores, cls_ids return boxes, scores, cls_ids, view_ids keep = np.ones(n, dtype=bool) for target_cls in self._parent_split_cls_ids(): class_indices = [i for i in range(n) if int(cls_ids[i]) == target_cls] if len(class_indices) <= 1: continue for a_idx in class_indices: area_a = self._box_area(boxes[a_idx]) small_indices = [ j for j in class_indices if j != a_idx and self._box_area(boxes[j]) < area_a and self._larger_involves_smaller( boxes, a_idx, j, involve_thresh ) ] if len(small_indices) < min_small_count: continue smalls_are_separate = True for i_pos in range(len(small_indices)): for j_pos in range(i_pos + 1, len(small_indices)): if self._boxes_involve_each_other( boxes, small_indices[i_pos], small_indices[j_pos], involve_thresh, ): smalls_are_separate = False break if not smalls_are_separate: break if smalls_are_separate: keep[a_idx] = False if view_ids is None: return boxes[keep], scores[keep], cls_ids[keep] return boxes[keep], scores[keep], cls_ids[keep], view_ids[keep] @staticmethod def _box_centroid(box: np.ndarray) -> tuple[float, float]: return ((box[0] + box[2]) / 2.0, (box[1] + box[3]) / 2.0) @staticmethod def _centroid_inside_box( boxes: np.ndarray, parent_idx: int, child_idx: int ) -> bool: parent = boxes[parent_idx] cx, cy = Miner._box_centroid(boxes[child_idx]) return parent[0] <= cx <= parent[2] and parent[1] <= cy <= parent[3] @staticmethod def _mostly_inside_parent( boxes: np.ndarray, parent_idx: int, child_idx: int, involve_thresh: float, ) -> bool: inter = Miner._box_intersection_area(boxes[parent_idx], boxes[child_idx]) if inter <= 0.0: return False area_child = Miner._box_area(boxes[child_idx]) if (inter / area_child) >= involve_thresh: return True child = boxes[child_idx] parent = boxes[parent_idx] cx = (child[0] + child[2]) / 2.0 cy = (child[1] + child[3]) / 2.0 return ( parent[0] <= cx <= parent[2] and parent[1] <= cy <= parent[3] ) def _cluster_indices_by_box_iou( self, boxes: np.ndarray, indices: list[int], iou_thresh: float, ) -> list[list[int]]: if not indices: return [] if len(indices) == 1: return [indices] parent_map = {idx: idx for idx in indices} def find(x: int) -> int: while parent_map[x] != x: parent_map[x] = parent_map[parent_map[x]] x = parent_map[x] return x def union(a: int, b: int) -> None: ra, rb = find(a), find(b) if ra != rb: parent_map[rb] = ra iou = self._compute_iou_matrix(boxes) for i_pos in range(len(indices)): for j_pos in range(i_pos + 1, len(indices)): a_idx = indices[i_pos] b_idx = indices[j_pos] if iou[a_idx, b_idx] > iou_thresh: union(a_idx, b_idx) groups: dict[int, list[int]] = {} for idx in indices: groups.setdefault(find(idx), []).append(idx) return list(groups.values()) def _cluster_indices_by_centroid_distance( self, boxes: np.ndarray, indices: list[int], max_dist: float, ) -> list[list[int]]: if not indices: return [] if len(indices) == 1: return [indices] parent_map = {idx: idx for idx in indices} def find(x: int) -> int: while parent_map[x] != x: parent_map[x] = parent_map[parent_map[x]] x = parent_map[x] return x def union(a: int, b: int) -> None: ra, rb = find(a), find(b) if ra != rb: parent_map[rb] = ra max_dist_sq = max_dist * max_dist centroids = [self._box_centroid(boxes[idx]) for idx in indices] for i_pos in range(len(indices)): cx_i, cy_i = centroids[i_pos] for j_pos in range(i_pos + 1, len(indices)): cx_j, cy_j = centroids[j_pos] dx = cx_i - cx_j dy = cy_i - cy_j if dx * dx + dy * dy <= max_dist_sq: union(indices[i_pos], indices[j_pos]) groups: dict[int, list[int]] = {} for idx in indices: groups.setdefault(find(idx), []).append(idx) return list(groups.values()) def _split_fire_parent_with_inner_clusters( self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray, view_ids: np.ndarray | None = None, inner_involve_thresh: float | None = None, cluster_iou_thresh: float | None = None, min_clusters: int | None = None, min_inner_conf: float | None = None, max_inner_area_ratio: float | None = None, peak_dist_ratio: float | None = None, peak_dist_min_px: float | None = None, peak_dist_max_px: float | None = None, peak_member_min_conf: float | None = None, ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray | None, np.ndarray]: if inner_involve_thresh is None: inner_involve_thresh = self.fire_split_inner_involve_thresh if cluster_iou_thresh is None: cluster_iou_thresh = self.fire_split_cluster_iou_thresh if min_clusters is None: min_clusters = self.fire_split_min_clusters if min_inner_conf is None: min_inner_conf = self.fire_split_min_inner_conf if max_inner_area_ratio is None: max_inner_area_ratio = self.fire_split_max_inner_area_ratio if peak_dist_ratio is None: peak_dist_ratio = self.fire_split_peak_dist_ratio if peak_dist_min_px is None: peak_dist_min_px = self.fire_split_peak_dist_min_px if peak_dist_max_px is None: peak_dist_max_px = self.fire_split_peak_dist_max_px if peak_member_min_conf is None: peak_member_min_conf = self.fire_split_peak_member_min_conf n = len(boxes) if n <= 1: protect = np.zeros(n, dtype=bool) return boxes, scores, cls_ids, view_ids, protect remove_indices: set[int] = set() split_boxes: list[np.ndarray] = [] split_scores: list[float] = [] split_cls: list[int] = [] split_views: list[int] = [] for target_cls in self._parent_split_cls_ids(): class_indices = [i for i in range(n) if int(cls_ids[i]) == target_cls] if len(class_indices) <= 1: continue class_indices_sorted = sorted( class_indices, key=lambda idx: self._box_area(boxes[idx]), reverse=True, ) for parent_idx in class_indices_sorted: if parent_idx in remove_indices: continue parent_area = self._box_area(boxes[parent_idx]) parent_scale = math.sqrt(parent_area) peak_dist = max( peak_dist_min_px, min(parent_scale * peak_dist_ratio, peak_dist_max_px), ) peak_indices = [ j for j in class_indices if j != parent_idx and j not in remove_indices and self._centroid_inside_box(boxes, parent_idx, j) ] max_inner_area = parent_area * max_inner_area_ratio small_inner_indices = [ j for j in peak_indices if self._box_area(boxes[j]) <= max_inner_area and self._mostly_inside_parent( boxes, parent_idx, j, inner_involve_thresh ) ] strong_clusters: list[list[int]] = [] member_indices: list[int] = [] use_small_split_outputs = False peak_clusters: list[list[int]] = [] peak_strong: list[list[int]] = [] peak_members = [parent_idx] + [ j for j in peak_indices if float(scores[j]) >= peak_member_min_conf ] if len(peak_members) >= min_clusters: peak_clusters = self._cluster_indices_by_centroid_distance( boxes, peak_members, peak_dist ) peak_strong = [ cluster for cluster in peak_clusters if max(float(scores[m]) for m in cluster) >= min_inner_conf ] if len(peak_strong) >= min_clusters: strong_clusters = peak_strong member_indices = peak_indices if len(strong_clusters) < min_clusters and small_inner_indices: iou_clusters = self._cluster_indices_by_box_iou( boxes, small_inner_indices, cluster_iou_thresh ) iou_strong = [ cluster for cluster in iou_clusters if max(float(scores[m]) for m in cluster) >= min_inner_conf ] if len(iou_strong) >= min_clusters: strong_clusters = iou_strong member_indices = small_inner_indices parent_blob_conf = 0.0 for cluster in peak_clusters: if parent_idx not in cluster: continue non_parent = [m for m in cluster if m != parent_idx] if non_parent: parent_blob_conf = max( parent_blob_conf, max(float(scores[m]) for m in non_parent), ) if ( len(strong_clusters) < min_clusters and parent_blob_conf >= self.fire_split_parent_blob_conf_thresh and len(small_inner_indices) >= min_clusters ): small_peak_clusters = self._cluster_indices_by_centroid_distance( boxes, small_inner_indices, peak_dist ) small_strong = [ cluster for cluster in small_peak_clusters if min( 1.0, max(float(scores[m]) for m in cluster) + self.fire_split_conf_boost, ) >= min_inner_conf ] if len(small_strong) >= min_clusters: cluster_best = [ cluster[int(np.argmax(scores[cluster]))] for cluster in small_strong ] smalls_are_separate = True for i_pos in range(len(cluster_best)): for j_pos in range(i_pos + 1, len(cluster_best)): if self._boxes_involve_each_other( boxes, cluster_best[i_pos], cluster_best[j_pos], self.fire_parent_involve_thresh, ): smalls_are_separate = False break if not smalls_are_separate: break if smalls_are_separate: strong_clusters = small_strong member_indices = peak_indices use_small_split_outputs = True if len(strong_clusters) < min_clusters: continue remove_indices.add(parent_idx) remove_indices.update(member_indices) for j in class_indices: if j in remove_indices: continue if self._boxes_involve_each_other( boxes, parent_idx, j, self.fire_parent_involve_thresh ): remove_indices.add(j) for cluster in strong_clusters: if use_small_split_outputs: best_member = cluster[int(np.argmax(scores[cluster]))] else: cluster_candidates = [ m for m in cluster if m != parent_idx ] or cluster best_member = cluster_candidates[ int(np.argmax(scores[cluster_candidates])) ] boosted_conf = float(scores[best_member]) member_area = self._box_area(boxes[best_member]) if ( use_small_split_outputs or member_area <= parent_area * max_inner_area_ratio ): boosted_conf = min( 1.0, boosted_conf + self.fire_split_conf_boost ) split_boxes.append(boxes[best_member].copy()) split_scores.append(boosted_conf) split_cls.append(target_cls) if view_ids is not None: split_views.append(int(view_ids[best_member])) if not remove_indices: protect = np.zeros(n, dtype=bool) return boxes, scores, cls_ids, view_ids, protect keep = np.array([i not in remove_indices for i in range(n)], dtype=bool) boxes = boxes[keep] scores = scores[keep] cls_ids = cls_ids[keep] if view_ids is not None: view_ids = view_ids[keep] protect = np.zeros(len(boxes), dtype=bool) if split_boxes: split_count = len(split_boxes) boxes = np.concatenate([boxes, np.stack(split_boxes, axis=0)], axis=0) scores = np.concatenate( [scores, np.asarray(split_scores, dtype=np.float32)], axis=0 ) cls_ids = np.concatenate( [cls_ids, np.asarray(split_cls, dtype=np.int32)], axis=0 ) if view_ids is not None: view_ids = np.concatenate( [view_ids, np.asarray(split_views, dtype=np.int32)], axis=0 ) else: view_ids = np.full(split_count, -1, dtype=np.int32) protect = np.concatenate( [protect, np.ones(split_count, dtype=bool)], axis=0 ) return boxes, scores, cls_ids, view_ids, protect def _remove_small_group_low_conf( self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray, min_count: int | None = None, conf_thresh: float | None = None, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: if min_count is None: min_count = self.small_group_min_count if conf_thresh is None: conf_thresh = self.small_group_low_conf_thresh n = len(boxes) if n == 0 or len(self.group_object_counts) != n: return boxes, scores, cls_ids keep = np.ones(n, dtype=bool) for i in range(n): if self.group_object_counts[i] < min_count and scores[i] < conf_thresh: keep[i] = False self.group_object_counts = [ count for count, k in zip(self.group_object_counts, keep) if k ] self.group_max_confidences = [ value for value, k in zip(self.group_max_confidences, keep) if k ] self.group_avg_confidences = [ value for value, k in zip(self.group_avg_confidences, keep) if k ] return boxes[keep], scores[keep], cls_ids[keep] def _apply_post_filters( self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray, orig_size: tuple[int, int], view_ids: np.ndarray | None = None, ) -> list[BoundingBox]: if len(boxes) == 0: return [] if view_ids is not None: boxes, scores, cls_ids, view_ids = self._remove_edge_low_conf( boxes, scores, cls_ids, orig_size, view_ids ) else: boxes, scores, cls_ids = self._remove_edge_low_conf( boxes, scores, cls_ids, orig_size ) if len(boxes) == 0: return [] split_protect: np.ndarray | None = None if view_ids is not None: boxes, scores, cls_ids, view_ids = ( self._remove_isolated_low_conf_same_class( boxes, scores, cls_ids, view_ids ) ) boxes, scores, cls_ids, view_ids = ( self._remove_overlapping_low_conf_same_class( boxes, scores, cls_ids, view_ids ) ) boxes, scores, cls_ids, view_ids, split_protect = ( self._split_fire_parent_with_inner_clusters( boxes, scores, cls_ids, view_ids ) ) if len(boxes) == 0: return [] boxes, scores, cls_ids, view_ids = ( self._remove_unconfirmed_single_view_low_conf( boxes, scores, cls_ids, view_ids, protect_mask=split_protect ) ) if len(boxes) == 0: return [] boxes, scores, cls_ids, view_ids = self._remove_spanning_parent_fire( boxes, scores, cls_ids, view_ids ) if len(boxes) == 0: return [] boxes, scores, cls_ids = self._group_overlapping_detections( boxes, scores, cls_ids, view_ids=view_ids, record_stats=False ) else: boxes, scores, cls_ids = self._remove_isolated_low_conf_same_class( boxes, scores, cls_ids ) boxes, scores, cls_ids = self._remove_overlapping_low_conf_same_class( boxes, scores, cls_ids ) boxes, scores, cls_ids, _, split_protect = ( self._split_fire_parent_with_inner_clusters( boxes, scores, cls_ids ) ) if len(boxes) == 0: return [] boxes, scores, cls_ids = self._remove_spanning_parent_fire( boxes, scores, cls_ids ) if len(boxes) == 0: return [] boxes, scores, cls_ids = self._group_overlapping_detections( boxes, scores, cls_ids, record_stats=True ) boxes, scores, cls_ids = self._remove_small_group_low_conf( boxes, scores, cls_ids ) boxes, scores, cls_ids = self._resolve_fire_and_ext_same_class_overlap( boxes, scores, cls_ids ) cls_smoke = self.class_names.index("smoke") if int(np.sum(cls_ids == cls_smoke)) >= self.smoke_cluster_merge_early_min_count: boxes, scores, cls_ids = self._merge_overlapping_smoke_clusters( boxes, scores, cls_ids ) boxes, scores, cls_ids = self._resolve_smoke_overlap( boxes, scores, cls_ids ) boxes, scores, cls_ids = self._merge_overlapping_smoke_clusters( boxes, scores, cls_ids ) return self._build_results(boxes, scores, cls_ids) @staticmethod def _bbox_to_array(box: BoundingBox) -> np.ndarray: return np.array([box.x1, box.y1, box.x2, box.y2], dtype=np.float32) def _smoke_base_band_box(self, smoke: BoundingBox) -> np.ndarray: smoke_h = smoke.y2 - smoke.y1 base_y1 = smoke.y1 + smoke_h * (1.0 - self.smoke_anchor_base_height_frac) return np.array( [float(smoke.x1), base_y1, float(smoke.x2), float(smoke.y2)], dtype=np.float32, ) def _fire_overlaps_smoke_base( self, fire_box: np.ndarray, smoke: BoundingBox ) -> bool: return self._boxes_involve_each_other_arrays( fire_box, self._smoke_base_band_box(smoke), self.smoke_anchor_corroborate_involve_thresh, ) def _fire_sits_under_smoke( self, fire_box: np.ndarray, smoke: BoundingBox ) -> bool: cx = (fire_box[0] + fire_box[2]) / 2.0 cy = (fire_box[1] + fire_box[3]) / 2.0 smoke_w = float(smoke.x2 - smoke.x1) smoke_cx = (smoke.x1 + smoke.x2) / 2.0 if abs(cx - smoke_cx) > smoke_w * 0.55: return False if cy < smoke.y1 + (smoke.y2 - smoke.y1) * 0.45: return False return True def _fire_centroid_in_smoke_base_band( self, fire_box: np.ndarray, smoke: BoundingBox ) -> bool: base = self._smoke_base_band_box(smoke) cx = (fire_box[0] + fire_box[2]) / 2.0 cy = (fire_box[1] + fire_box[3]) / 2.0 return ( base[0] <= cx <= base[2] and base[1] <= cy <= base[3] ) def _corroborating_merged_tta_fire_boxes( self, fire_box: np.ndarray, merged_boxes: np.ndarray, merged_cls: np.ndarray, smoke: BoundingBox, ) -> list[np.ndarray]: cls_fire = self.class_names.index("fire") involve_thresh = self.smoke_anchor_corroborate_involve_thresh probe_area = self._box_area(fire_box) corroborators: list[np.ndarray] = [] for i in range(len(merged_boxes)): if int(merged_cls[i]) != cls_fire: continue merged_box = merged_boxes[i] if not self._boxes_involve_each_other_arrays( fire_box, merged_box, involve_thresh ): continue if not self._fire_centroid_in_smoke_base_band(merged_box, smoke): continue merged_area = self._box_area(merged_box) if merged_area <= 0.0 or probe_area <= 0.0: continue area_ratio = probe_area / merged_area if area_ratio > self.smoke_anchor_corroborate_area_ratio_max: continue if area_ratio < 1.0 / self.smoke_anchor_corroborate_area_ratio_max: continue corroborators.append(merged_box.copy()) return corroborators def _snap_smoke_anchor_probe_to_corroborator( self, fire_box: np.ndarray, merged_boxes: np.ndarray, merged_cls: np.ndarray, smoke: BoundingBox, crop_size: tuple[int, int], ) -> np.ndarray | None: corroborators = self._corroborating_merged_tta_fire_boxes( fire_box, merged_boxes, merged_cls, smoke ) if not corroborators: return None corroborators.sort(key=self._box_area) for corr in corroborators: if self._passes_smoke_anchor_probe_fire_max_size( corr, smoke, crop_size ): return corr.copy() return None def _corroborated_by_merged_tta_fire( self, fire_box: np.ndarray, merged_boxes: np.ndarray, merged_cls: np.ndarray, smoke: BoundingBox, ) -> bool: return bool( self._corroborating_merged_tta_fire_boxes( fire_box, merged_boxes, merged_cls, smoke ) ) def _refine_smoke_anchor_probe_fire_box( self, probe_box: np.ndarray, image: np.ndarray, smoke: BoundingBox, merged_boxes: np.ndarray, merged_cls: np.ndarray, crop_size: tuple[int, int], ) -> np.ndarray: corroborators = self._corroborating_merged_tta_fire_boxes( probe_box, merged_boxes, merged_cls, smoke ) if not corroborators: return probe_box def _passes_refined(box: np.ndarray) -> bool: return ( self._passes_smoke_anchor_probe_fire_max_size(box, smoke, crop_size) and self._passes_smoke_anchor_probe_fire_bottom_dist(box, smoke) and self._passes_smoke_anchor_probe_fire_color(image, box) ) probe_area = self._box_area(probe_box) if probe_area <= 0.0: return probe_box expanders: list[np.ndarray] = [] for corr in corroborators: if not self._passes_smoke_anchor_probe_fire_max_size( corr, smoke, crop_size ): continue for candidate in ( corr, np.array( [ min(probe_box[0], corr[0]), min(probe_box[1], corr[1]), max(probe_box[2], corr[2]), max(probe_box[3], corr[3]), ], dtype=np.float32, ), ): if not _passes_refined(candidate): continue if ( candidate[0] > probe_box[0] + 1.0 and candidate[1] > probe_box[1] + 1.0 ): continue inter = self._box_intersection_area(probe_box, candidate) if inter / probe_area < 0.95: continue expanders.append(candidate) if expanders: refined = min( expanders, key=lambda box: (box[0], box[2] - box[0], box[1]), ) else: refined = probe_box pad = self.smoke_anchor_corroborated_refine_pad_px if pad > 0.0 and corroborators: padded = refined.copy() padded[0] = max(0.0, padded[0] - pad) padded[1] = max(0.0, padded[1] - pad) if _passes_refined(padded): refined = padded return refined def _smoke_anchor_min_fire_size( self, corroborated: bool, smoke_conf: float, under_base: bool = False, ) -> tuple[float, float]: if smoke_conf >= self.smoke_anchor_strong_smoke_conf and ( corroborated or under_base ): return ( self.smoke_anchor_corroborated_min_fire_width_px, self.smoke_anchor_corroborated_min_fire_height_px, ) return ( self.smoke_anchor_min_fire_width_px, self.smoke_anchor_min_fire_height_px, ) def _passes_smoke_anchor_fire_size( self, fire_box: np.ndarray, corroborated: bool, smoke_conf: float, under_base: bool = False, ) -> bool: box_w = float(fire_box[2] - fire_box[0]) box_h = float(fire_box[3] - fire_box[1]) min_w, min_h = self._smoke_anchor_min_fire_size( corroborated, smoke_conf, under_base ) return box_w >= min_w and box_h >= min_h def _passes_smoke_anchor_probe_fire_max_size( self, fire_box: np.ndarray, smoke: BoundingBox, crop_size: tuple[int, int], ) -> bool: box_w = float(fire_box[2] - fire_box[0]) box_h = float(fire_box[3] - fire_box[1]) if box_w > self.smoke_anchor_max_probe_fire_width_px: return False if box_h > self.smoke_anchor_max_probe_fire_height_px: return False smoke_area = max( 1.0, float((smoke.x2 - smoke.x1) * (smoke.y2 - smoke.y1)) ) if self._box_area(fire_box) / smoke_area > self.smoke_anchor_max_probe_fire_area_ratio: return False crop_w, crop_h = crop_size if crop_w > 0 and box_w / float(crop_w) > self.smoke_anchor_max_probe_crop_width_ratio: return False if crop_h > 0 and box_h / float(crop_h) > self.smoke_anchor_max_probe_crop_height_ratio: return False return True def _passes_smoke_anchor_probe_fire_bottom_dist( self, fire_box: np.ndarray, smoke: BoundingBox ) -> bool: smoke_h = max(1.0, float(smoke.y2 - smoke.y1)) bottom_gap = float(smoke.y2) - float(fire_box[3]) return bottom_gap / smoke_h <= self.smoke_anchor_max_probe_fire_bottom_dist_frac def _passes_smoke_anchor_probe_fire_color( self, image: np.ndarray, fire_box: np.ndarray ) -> bool: h, w = image.shape[:2] x1 = max(0, int(math.floor(fire_box[0]))) y1 = max(0, int(math.floor(fire_box[1]))) x2 = min(w, int(math.ceil(fire_box[2]))) y2 = min(h, int(math.ceil(fire_box[3]))) if x2 <= x1 or y2 <= y1: return False roi = image[y1:y2, x1:x2] if roi.size == 0: return False blue = roi[:, :, 0].astype(np.float32) green = roi[:, :, 1].astype(np.float32) red = roi[:, :, 2].astype(np.float32) mean_r = float(np.mean(red)) max_rgb = float( max(np.max(red), np.max(green), np.max(blue)) ) bright_frac = float( np.mean(np.max(roi, axis=2) >= self.smoke_anchor_probe_bright_thresh) ) if ( max_rgb >= self.smoke_anchor_probe_sparse_hot_min_max_rgb and bright_frac >= self.smoke_anchor_probe_sparse_hot_min_bright_frac ): return True if mean_r < self.smoke_anchor_probe_min_mean_r: return False if max_rgb < self.smoke_anchor_probe_min_max_rgb: return False warm_mask = (red > green + 10.0) & (red > blue + 10.0) warm_frac = float(np.mean(warm_mask)) r_minus_g = mean_r - float(np.mean(green)) if warm_frac >= self.smoke_anchor_probe_min_warm_frac: if ( max_rgb >= self.smoke_anchor_probe_warm_min_max_rgb or mean_r >= self.smoke_anchor_probe_warm_min_mean_r or warm_frac >= self.smoke_anchor_probe_strong_warm_frac ): return True if ( bright_frac >= self.smoke_anchor_probe_min_bright_frac and r_minus_g >= self.smoke_anchor_probe_min_r_minus_g ): return True if ( max_rgb >= self.smoke_anchor_probe_min_hot_max_rgb and r_minus_g >= self.smoke_anchor_probe_min_r_minus_g ): return True return False def _fire_ext_red_distribution( self, image: np.ndarray, ext_box: np.ndarray ) -> tuple[float, float, float]: h, w = image.shape[:2] x1 = max(0, int(math.floor(float(ext_box[0])))) y1 = max(0, int(math.floor(float(ext_box[1])))) x2 = min(w, int(math.ceil(float(ext_box[2])))) y2 = min(h, int(math.ceil(float(ext_box[3])))) if x2 <= x1 or y2 <= y1: return 0.0, 0.0, 0.0 roi = image[y1:y2, x1:x2] if roi.size == 0: return 0.0, 0.0, 0.0 blue = roi[:, :, 0].astype(np.float32) green = roi[:, :, 1].astype(np.float32) red = roi[:, :, 2].astype(np.float32) red_dom = (red > green + 10.0) & (red > blue + 10.0) return ( float(np.mean(red)), float(np.mean(red_dom)), float(np.mean(red - green)), ) def _passes_fire_ext_red_color( self, image: np.ndarray, ext_box: np.ndarray ) -> bool: mean_r, red_dom, r_minus_g = self._fire_ext_red_distribution( image, ext_box ) if red_dom >= self.fire_ext_min_red_dom_frac: return True if ( r_minus_g >= self.fire_ext_min_r_minus_g and mean_r >= self.fire_ext_color_min_mean_r ): return True return False def _filter_fire_ext_by_red_color( self, image: np.ndarray, results: list[BoundingBox] ) -> list[BoundingBox]: cls_fire_ext = self.class_names.index("fire extinguisher") max_conf = self.fire_ext_color_filter_max_conf filtered: list[BoundingBox] = [] for box in results: if ( box.cls_id == cls_fire_ext and box.conf <= max_conf + 1e-4 and not self._passes_fire_ext_red_color( image, self._bbox_to_array(box) ) ): continue filtered.append(box) return filtered def _filter_probe_fires_by_color( self, image: np.ndarray, results: list[BoundingBox] ) -> list[BoundingBox]: cls_fire = self.class_names.index("fire") floor_conf = self.smoke_anchor_fire_floor_conf filtered: list[BoundingBox] = [] for box in results: if ( box.cls_id == cls_fire and box.conf <= floor_conf + 1e-4 and not self._passes_smoke_anchor_probe_fire_color( image, self._bbox_to_array(box) ) ): continue filtered.append(box) return filtered def _smoke_anchor_fire_is_duplicate( self, fire_box: np.ndarray, others: list[BoundingBox], cls_fire: int, involve_thresh: float | None = None, ) -> bool: if involve_thresh is None: involve_thresh = self.smoke_anchor_corroborate_involve_thresh for fire in others: if fire.cls_id != cls_fire: continue if self._boxes_involve_each_other_arrays( self._bbox_to_array(fire), fire_box, involve_thresh ): return True return False def _collect_smoke_anchor_probe_fires( self, image: np.ndarray, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray, smoke: BoundingBox, merged_boxes: np.ndarray, merged_cls: np.ndarray, existing_fires: list[BoundingBox], cls_fire: int, crop_size: tuple[int, int], ) -> list[tuple[np.ndarray, float]]: smoke_area = max( 1.0, float((smoke.x2 - smoke.x1) * (smoke.y2 - smoke.y1)) ) candidates: list[tuple[float, np.ndarray]] = [] for i in range(len(boxes)): if int(cls_ids[i]) != cls_fire: continue fire_box = boxes[i] corroborated = self._corroborated_by_merged_tta_fire( fire_box, merged_boxes, merged_cls, smoke ) under_smoke = self._fire_sits_under_smoke(fire_box, smoke) in_base = self._fire_centroid_in_smoke_base_band(fire_box, smoke) under_base = under_smoke and in_base if not corroborated and not under_base: continue if not self._passes_smoke_anchor_fire_size( fire_box, corroborated, smoke.conf, under_base ): continue if not self._passes_smoke_anchor_probe_fire_max_size( fire_box, smoke, crop_size ): if corroborated: fire_box = self._snap_smoke_anchor_probe_to_corroborator( fire_box, merged_boxes, merged_cls, smoke, crop_size ) if fire_box is None: continue else: continue if not self._passes_smoke_anchor_probe_fire_bottom_dist(fire_box, smoke): continue if not self._passes_smoke_anchor_probe_fire_color(image, fire_box): continue if self._box_area(fire_box) > smoke_area * self.smoke_anchor_max_fire_area_ratio: continue candidates.append((float(scores[i]), fire_box.copy())) candidates.sort(key=lambda item: item[0], reverse=True) selected: list[tuple[np.ndarray, float]] = [] separate_thresh = self.smoke_anchor_separate_fire_involve_thresh max_fires = self.smoke_anchor_max_fires_per_smoke floor_conf = self.smoke_anchor_fire_floor_conf for score, fire_box in candidates: corroborated = self._corroborated_by_merged_tta_fire( fire_box, merged_boxes, merged_cls, smoke ) under_smoke = self._fire_sits_under_smoke(fire_box, smoke) in_base = self._fire_centroid_in_smoke_base_band(fire_box, smoke) min_conf = floor_conf if corroborated and under_smoke and in_base: min_conf = self.smoke_anchor_corroborated_probe_min_conf elif not corroborated and under_smoke and in_base: min_conf = self.smoke_anchor_probe_rescue_min_conf if score < min_conf: continue if len(selected) >= max_fires: break if self._smoke_anchor_fire_is_duplicate( fire_box, existing_fires, cls_fire ): continue overlaps_selected = False for chosen_box, _ in selected: if self._boxes_involve_each_other_arrays( fire_box, chosen_box, separate_thresh ): overlaps_selected = True break if overlaps_selected: continue selected.append((fire_box, score)) return selected def _smoke_anchor_crop_region( self, smoke: BoundingBox, orig_size: tuple[int, int] ) -> tuple[int, int, int, int]: orig_w, orig_h = orig_size smoke_w = float(smoke.x2 - smoke.x1) smoke_h = float(smoke.y2 - smoke.y1) pad_x = smoke_w * self.smoke_anchor_crop_width_pad_ratio y1 = smoke.y1 + smoke_h * (1.0 - self.smoke_anchor_base_height_frac) y2 = min(orig_h, float(smoke.y2) + self.smoke_anchor_crop_down_pad_px) x1 = max(0, int(math.floor(smoke.x1 - pad_x))) x2 = min(orig_w, int(math.ceil(smoke.x2 + pad_x))) y1i = max(0, int(math.floor(y1))) y2i = min(orig_h, max(y1i + 1, int(math.ceil(y2)))) x2 = max(x1 + 1, x2) return x1, x2, y1i, y2i def _smoke_anchor_probe_target_smoke( self, smoke: BoundingBox, smoke_seeds: list[BoundingBox] | None, ) -> BoundingBox: if not smoke_seeds: return smoke cls_smoke = self.class_names.index("smoke") smoke_arr = self._bbox_to_array(smoke) smoke_w = float(smoke.x2 - smoke.x1) best_seed: BoundingBox | None = None for seed in smoke_seeds: if seed.cls_id != cls_smoke: continue if seed.conf < self.smoke_anchor_min_smoke_conf: continue seed_arr = self._bbox_to_array(seed) if not self._boxes_involve_each_other_arrays( smoke_arr, seed_arr, self.smoke_expand_corroborate_involve_thresh ): continue seed_w = float(seed.x2 - seed.x1) if seed_w >= smoke_w * 0.95: continue if best_seed is None or seed.conf > best_seed.conf: best_seed = seed if best_seed is None: return smoke seed_w = float(best_seed.x2 - best_seed.x1) if smoke_w > seed_w * self.smoke_anchor_expanded_width_ratio: return best_seed return smoke def _probe_smoke_anchored_fire( self, image: np.ndarray, results: list[BoundingBox], orig_size: tuple[int, int], merged_boxes: np.ndarray, merged_cls: np.ndarray, smoke_seeds: list[BoundingBox] | None = None, ) -> list[BoundingBox]: if not self.use_tta: return results cls_fire = self.class_names.index("fire") cls_smoke = self.class_names.index("smoke") floor_conf = self.smoke_anchor_fire_floor_conf added: list[BoundingBox] = [] smoke_candidates = sorted( ( box for box in results if box.cls_id == cls_smoke and box.conf >= self.smoke_anchor_min_smoke_conf ), key=lambda box: box.conf, reverse=True, )[: self.smoke_anchor_max_probe_smokes] if not smoke_candidates: return results for smoke in smoke_candidates: smoke_arr = self._bbox_to_array(smoke) fire_on_smoke_base = False for fire in results: if fire.cls_id != cls_fire or fire.conf < self.fire_min_conf: continue if self._fire_overlaps_smoke_base( self._bbox_to_array(fire), smoke ): fire_on_smoke_base = True break if fire_on_smoke_base: continue anchor_smoke = smoke x1, x2, y1, y2 = self._smoke_anchor_crop_region(anchor_smoke, orig_size) crop = image[y1:y2, x1:x2] if crop.size == 0: continue crop_size = (crop.shape[1], crop.shape[0]) boxes, scores, cls_ids = self._infer_view_arrays( crop, crop_size, flip_code=None ) if len(boxes) == 0: continue boxes = boxes.copy() boxes[:, [0, 2]] += x1 boxes[:, [1, 3]] += y1 boxes = self._clip_boxes(boxes, orig_size) probe_fires = self._collect_smoke_anchor_probe_fires( image, boxes, scores, cls_ids, anchor_smoke, merged_boxes, merged_cls, results + added, cls_fire, crop_size, ) if not probe_fires: seed_smoke = self._smoke_anchor_probe_target_smoke(smoke, smoke_seeds) if seed_smoke is not smoke: x1, x2, y1, y2 = self._smoke_anchor_crop_region( seed_smoke, orig_size ) crop = image[y1:y2, x1:x2] if crop.size == 0: continue crop_size = (crop.shape[1], crop.shape[0]) boxes, scores, cls_ids = self._infer_view_arrays( crop, crop_size, flip_code=None ) if len(boxes) == 0: continue boxes = boxes.copy() boxes[:, [0, 2]] += x1 boxes[:, [1, 3]] += y1 boxes = self._clip_boxes(boxes, orig_size) probe_fires = self._collect_smoke_anchor_probe_fires( image, boxes, scores, cls_ids, seed_smoke, merged_boxes, merged_cls, results + added, cls_fire, crop_size, ) anchor_smoke = seed_smoke if not probe_fires: continue for probe_box, probe_raw_score in probe_fires: probe_box = self._refine_smoke_anchor_probe_fire_box( probe_box, image, anchor_smoke, merged_boxes, merged_cls, crop_size, ) probe_conf = min(1.0, max(probe_raw_score, floor_conf)) added.append( BoundingBox( x1=int(math.floor(probe_box[0])), y1=int(math.floor(probe_box[1])), x2=int(math.ceil(probe_box[2])), y2=int(math.ceil(probe_box[3])), cls_id=cls_fire, conf=probe_conf, ) ) if len(added) > self.smoke_anchor_max_crop_fires: added.sort(key=lambda box: box.conf, reverse=True) added = added[: self.smoke_anchor_max_crop_fires] return results + added def _ext_probe_union_box( self, boxes: np.ndarray, indices: list[int] ) -> np.ndarray: cluster_boxes = boxes[indices] return np.array( [ float(np.min(cluster_boxes[:, 0])), float(np.min(cluster_boxes[:, 1])), float(np.max(cluster_boxes[:, 2])), float(np.max(cluster_boxes[:, 3])), ], dtype=np.float32, ) def _ext_probe_seed_on_border( self, seed_box: np.ndarray, orig_size: tuple[int, int] ) -> bool: orig_w, orig_h = orig_size tol = self.ext_probe_edge_tol return bool( seed_box[0] <= tol or seed_box[1] <= tol or seed_box[2] >= orig_w - 1 - tol or seed_box[3] >= orig_h - 1 - tol ) def _ext_probe_singleton_too_close_to_border( self, seed_box: np.ndarray, orig_size: tuple[int, int] ) -> bool: orig_w, orig_h = orig_size margin = self.ext_probe_singleton_min_edge_margin_px return bool( seed_box[0] < margin or seed_box[1] < margin or seed_box[2] > orig_w - 1 - margin or seed_box[3] > orig_h - 1 - margin ) def _ext_probe_overlaps_strong_fire( self, seed_box: np.ndarray, results: list[BoundingBox] ) -> bool: cls_fire = self.class_names.index("fire") for box in results: if box.cls_id != cls_fire or box.conf < self.fire_min_conf: continue if self._boxes_involve_each_other_arrays( seed_box, self._bbox_to_array(box), self.ext_probe_fire_exclude_involve_thresh, ): return True return False def _ext_probe_near_any_fire( self, seed_box: np.ndarray, results: list[BoundingBox] ) -> bool: cls_fire = self.class_names.index("fire") max_gap = self.ext_probe_near_fire_gap_px for box in results: if box.cls_id != cls_fire: continue fire_arr = self._bbox_to_array(box) if self._box_axis_gap(seed_box, fire_arr, 0) > max_gap: continue if self._box_axis_gap(seed_box, fire_arr, 1) <= max_gap: return True return False def _ext_probe_is_duplicate( self, ext_box: np.ndarray, results: list[BoundingBox], cls_fire_ext: int, min_conf: float = 0.0, ) -> bool: for box in results: if box.cls_id != cls_fire_ext or box.conf < min_conf: continue if self._boxes_involve_each_other_arrays( ext_box, self._bbox_to_array(box), self.ext_probe_duplicate_involve_thresh, ): return True return False def _passes_ext_probe_size( self, ext_box: np.ndarray, crop_score: float = 0.0 ) -> bool: box_w = float(ext_box[2] - ext_box[0]) box_h = float(ext_box[3] - ext_box[1]) min_w = self.ext_probe_min_width_px if crop_score >= self.ext_probe_crop_rescue_min_conf: min_w = min(min_w, self.ext_probe_crop_rescue_min_width_px) return ( box_w >= min_w and box_h >= self.ext_probe_min_height_px and box_w <= self.ext_probe_max_width_px and box_h <= self.ext_probe_max_height_px ) def _passes_ext_probe_singleton_size(self, ext_box: np.ndarray) -> bool: box_w = float(ext_box[2] - ext_box[0]) box_h = float(ext_box[3] - ext_box[1]) return ( box_w >= self.ext_probe_min_width_px and box_h >= self.ext_probe_min_height_px and box_w <= self.ext_probe_singleton_max_width_px and box_h <= self.ext_probe_singleton_max_height_px and self._box_area(ext_box) <= self.ext_probe_singleton_max_area ) def _ext_probe_validate_seed( self, seed_box: np.ndarray, best_score: float, min_conf: float, results: list[BoundingBox], orig_size: tuple[int, int], cls_fire_ext: int, ) -> bool: if best_score < min_conf: return False if best_score >= self.fire_ext_min_conf: return False if not self._passes_ext_probe_size(seed_box): return False if self._ext_probe_seed_on_border(seed_box, orig_size): return False if self._ext_probe_overlaps_strong_fire(seed_box, results): return False if self._ext_probe_near_any_fire(seed_box, results): return False if self._ext_probe_is_duplicate( seed_box, results, cls_fire_ext, min_conf=self.fire_ext_min_conf ): return False return True def _ext_probe_crop_region( self, seed_box: np.ndarray, orig_size: tuple[int, int] ) -> tuple[int, int, int, int]: orig_w, orig_h = orig_size seed_w = float(seed_box[2] - seed_box[0]) seed_h = float(seed_box[3] - seed_box[1]) pad_x = seed_w * self.ext_probe_crop_pad_ratio pad_y = seed_h * self.ext_probe_crop_pad_ratio x1 = max(0, int(math.floor(seed_box[0] - pad_x))) x2 = min(orig_w, int(math.ceil(seed_box[2] + pad_x))) y1 = max(0, int(math.floor(seed_box[1] - pad_y))) y2 = min(orig_h, int(math.ceil(seed_box[3] + pad_y))) min_crop = self.ext_probe_min_crop_size_px crop_w = float(x2 - x1) crop_h = float(y2 - y1) if crop_w < min_crop: cx = 0.5 * (float(seed_box[0]) + float(seed_box[2])) half = 0.5 * min_crop x1 = max(0, int(math.floor(cx - half))) x2 = min(orig_w, int(math.ceil(cx + half))) if crop_h < min_crop: cy = 0.5 * (float(seed_box[1]) + float(seed_box[3])) half = 0.5 * min_crop y1 = max(0, int(math.floor(cy - half))) y2 = min(orig_h, int(math.ceil(cy + half))) x2 = max(x1 + 1, x2) y2 = max(y1 + 1, y2) return x1, x2, y1, y2 def _collect_merged_ext_probe_seeds( self, merged_boxes: np.ndarray, merged_scores: np.ndarray, merged_cls: np.ndarray, results: list[BoundingBox], orig_size: tuple[int, int], ) -> list[tuple[np.ndarray, float, list[int]]]: cls_fire_ext = self.class_names.index("fire extinguisher") ext_indices = [ i for i in range(len(merged_boxes)) if int(merged_cls[i]) == cls_fire_ext ] if len(ext_indices) == 0: return [] clusters = self._cluster_indices_by_centroid_distance( merged_boxes, ext_indices, self.ext_probe_cluster_centroid_dist_px, ) seeds: list[tuple[np.ndarray, float, list[int]]] = [] for cluster in clusters: best_idx = cluster[int(np.argmax(merged_scores[cluster]))] best_score = float(merged_scores[best_idx]) if len(cluster) >= self.ext_probe_min_cluster_count: if best_score < self.ext_probe_min_cluster_best_conf: continue seed_box = self._ext_probe_union_box(merged_boxes, cluster) min_conf = self.ext_probe_min_cluster_best_conf elif len(cluster) == 1: seed_box = merged_boxes[best_idx].copy() if not self._passes_ext_probe_singleton_size(seed_box): continue if self._ext_probe_singleton_too_close_to_border(seed_box, orig_size): continue min_conf = self.ext_probe_singleton_min_conf else: continue if not self._ext_probe_validate_seed( seed_box, best_score, min_conf, results, orig_size, cls_fire_ext ): continue seeds.append((seed_box, best_score, cluster)) seeds.sort(key=lambda item: item[1], reverse=True) return seeds def _collect_ext_probe_crop_candidates( self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray, seed_box: np.ndarray, cls_fire_ext: int, ) -> list[tuple[np.ndarray, float]]: candidates: list[tuple[float, np.ndarray]] = [] for i in range(len(boxes)): if int(cls_ids[i]) != cls_fire_ext: continue ext_box = boxes[i] if not self._boxes_involve_each_other_arrays( seed_box, ext_box, self.ext_probe_duplicate_involve_thresh ): continue if not self._passes_ext_probe_size(ext_box, float(scores[i])): continue candidates.append((float(scores[i]), ext_box.copy())) candidates.sort(key=lambda item: item[0], reverse=True) return [(box, score) for score, box in candidates] def _refine_ext_probe_box( self, seed_box: np.ndarray, crop_candidates: list[tuple[np.ndarray, float]], ) -> tuple[np.ndarray, float]: if not crop_candidates: return seed_box.copy(), 0.0 crop_box, crop_score = crop_candidates[0] if crop_score < self.ext_probe_rescue_min_conf: return seed_box.copy(), 0.0 refined = np.array( [ min(float(seed_box[0]), float(crop_box[0])), min(float(seed_box[1]), float(crop_box[1])), max(float(seed_box[2]), float(crop_box[2])), max(float(seed_box[3]), float(crop_box[3])), ], dtype=np.float32, ) if not self._passes_ext_probe_size(refined, crop_score): return crop_box.copy(), crop_score return refined, crop_score def _probe_merged_anchored_fire_ext( self, image: np.ndarray, results: list[BoundingBox], orig_size: tuple[int, int], merged_boxes: np.ndarray, merged_scores: np.ndarray, merged_cls: np.ndarray, ) -> list[BoundingBox]: if not self.use_tta: return results cls_fire_ext = self.class_names.index("fire extinguisher") floor_conf = self.fire_ext_min_conf seeds = self._collect_merged_ext_probe_seeds( merged_boxes, merged_scores, merged_cls, results, orig_size ) if not seeds: return results added: list[BoundingBox] = [] for seed_box, seed_score, _cluster in seeds: x1, x2, y1, y2 = self._ext_probe_crop_region(seed_box, orig_size) crop = image[y1:y2, x1:x2] crop_score = 0.0 probe_box = seed_box.copy() if crop.size > 0: crop_size = (crop.shape[1], crop.shape[0]) boxes, scores, cls_ids = self._infer_view_arrays( crop, crop_size, flip_code=None ) if len(boxes) > 0: boxes = boxes.copy() boxes[:, [0, 2]] += x1 boxes[:, [1, 3]] += y1 boxes = self._clip_boxes(boxes, orig_size) crop_candidates = self._collect_ext_probe_crop_candidates( boxes, scores, cls_ids, seed_box, cls_fire_ext ) probe_box, crop_score = self._refine_ext_probe_box( seed_box, crop_candidates ) if not self._passes_ext_probe_size(probe_box, crop_score): continue if self._ext_probe_is_duplicate( probe_box, results + added, cls_fire_ext, min_conf=self.fire_ext_min_conf, ): continue probe_conf = min( 1.0, max(seed_score, crop_score, floor_conf) ) added.append( BoundingBox( x1=int(math.floor(probe_box[0])), y1=int(math.floor(probe_box[1])), x2=int(math.ceil(probe_box[2])), y2=int(math.ceil(probe_box[3])), cls_id=cls_fire_ext, conf=probe_conf, ) ) if len(added) >= self.ext_probe_max_per_frame: break if not added: return results kept: list[BoundingBox] = [] for box in results: if box.cls_id == cls_fire_ext and box.conf < floor_conf: if any( self._boxes_involve_each_other_arrays( self._bbox_to_array(box), self._bbox_to_array(added_box), self.ext_probe_duplicate_involve_thresh, ) for added_box in added ): continue kept.append(box) return kept + added def _smoke_expansion_max_deltas( self, orig_size: tuple[int, int] ) -> tuple[float, float]: orig_w, orig_h = orig_size return ( float(orig_w) * self.smoke_expand_max_frame_ratio, float(orig_h) * self.smoke_expand_max_frame_ratio, ) def _box_axis_gap( self, box_a: np.ndarray, box_b: np.ndarray, axis: int ) -> float: start_a = float(box_a[axis]) end_a = float(box_a[axis + 2]) start_b = float(box_b[axis]) end_b = float(box_b[axis + 2]) if end_a >= start_b and end_b >= start_a: return 0.0 return max(start_b - end_a, start_a - end_b) def _smoke_probe_within_expansion_reach( self, seed_arr: np.ndarray, probe_box: np.ndarray, orig_size: tuple[int, int], ) -> bool: max_dx, max_dy = self._smoke_expansion_max_deltas(orig_size) if self._box_axis_gap(seed_arr, probe_box, 0) > max_dx: return False if self._box_axis_gap(seed_arr, probe_box, 1) > max_dy: return False return True def _clamp_smoke_expansion_union( self, union_box: np.ndarray, seed_arr: np.ndarray, orig_size: tuple[int, int], ) -> np.ndarray: max_dx, max_dy = self._smoke_expansion_max_deltas(orig_size) union_box = union_box.copy() union_box[0] = max(float(union_box[0]), float(seed_arr[0]) - max_dx) union_box[2] = min(float(union_box[2]), float(seed_arr[2]) + max_dx) union_box[1] = max(float(union_box[1]), float(seed_arr[1]) - max_dy) union_box[3] = min(float(union_box[3]), float(seed_arr[3]) + max_dy) return union_box def _smoke_crop_horizontal_reach_limit(self, seed_arr: np.ndarray) -> float: seed_w = max(1.0, float(seed_arr[2]) - float(seed_arr[0])) return max( 32.0, seed_w * self.smoke_expand_band_max_horizontal_gap_ratio, ) def _clip_smoke_crop_candidate_for_union( self, candidate_box: np.ndarray, seed_arr: np.ndarray, ) -> np.ndarray: if self._probe_qualifies_for_upward_low_conf(candidate_box, seed_arr): return candidate_box.copy() max_h_gap = self._smoke_crop_horizontal_reach_limit(seed_arr) clipped = candidate_box.copy() clipped[0] = max(float(clipped[0]), float(seed_arr[0]) - max_h_gap) clipped[2] = min(float(clipped[2]), float(seed_arr[2]) + max_h_gap) if float(clipped[2]) <= float(clipped[0]): clipped[0] = float(seed_arr[0]) clipped[2] = float(seed_arr[2]) return clipped def _smoke_expansion_crop_region( self, smoke: BoundingBox, orig_size: tuple[int, int] ) -> tuple[int, int, int, int]: orig_w, orig_h = orig_size max_dx, max_dy = self._smoke_expansion_max_deltas(orig_size) x1 = max(0, int(math.floor(float(smoke.x1) - max_dx))) x2 = min(orig_w, int(math.ceil(float(smoke.x2) + max_dx))) y1 = max(0, int(math.floor(float(smoke.y1) - max_dy))) y2 = min( orig_h, max(y1 + 1, int(math.ceil(float(smoke.y2) + max_dy))), ) x2 = max(x1 + 1, min(orig_w, x2)) y2 = max(y1 + 1, min(orig_h, y2)) return x1, x2, y1, y2 def _smoke_expansion_uses_wide_crop( self, smoke: BoundingBox, orig_size: tuple[int, int] ) -> bool: orig_w, orig_h = orig_size smoke_w = float(smoke.x2 - smoke.x1) return ( float(smoke.y2) < orig_h * self.smoke_expand_upper_plume_max_y_ratio and smoke_w < orig_w * self.smoke_expand_wide_seed_max_width_ratio ) def _smoke_expansion_crop_regions( self, smoke: BoundingBox, orig_size: tuple[int, int] ) -> list[tuple[int, int, int, int]]: primary = self._smoke_expansion_crop_region(smoke, orig_size) regions = [primary] if not self._smoke_expansion_uses_wide_crop(smoke, orig_size): return regions orig_w, orig_h = orig_size px1, px2, py1, py2 = primary max_dy = self._smoke_expansion_max_deltas(orig_size)[1] smoke_h = max(1.0, float(smoke.y2 - smoke.y1)) slice_y1 = max( py1, int(math.floor(float(smoke.y1) - min(smoke_h * 0.15, max_dy))), ) slice_y2 = min( py2, max(slice_y1 + 1, int(math.ceil(float(smoke.y2) + max_dy))), ) left_x2 = max(px1 + 1, min(px2, int(math.ceil(float(smoke.x2))))) right_x1 = min(px2 - 1, max(px1, int(math.floor(float(smoke.x1))))) min_slice_w = max(40, int(orig_w * 0.08)) if left_x2 - px1 >= min_slice_w: regions.append((px1, left_x2, slice_y1, slice_y2)) if px2 - right_x1 >= min_slice_w: regions.append((right_x1, px2, slice_y1, slice_y2)) return regions def _smoke_corroborates_upward_diagonal_plume( self, probe_box: np.ndarray, seed_arr: np.ndarray, crop_x1: int, crop_x2: int, ) -> bool: if not self._probe_extends_smoke_upward(probe_box, seed_arr): return False smoke_h = max(1.0, float(seed_arr[3]) - seed_arr[1]) smoke_w = max(1.0, float(seed_arr[2]) - seed_arr[0]) if float(probe_box[3]) < float(seed_arr[1]) - smoke_h * 0.15: return False cy = (float(probe_box[1]) + float(probe_box[3])) / 2.0 if cy > float(seed_arr[1]) + smoke_h * 0.35: return False gap_x = self._box_axis_gap(probe_box, seed_arr, 0) max_gap = max( 32.0, smoke_w * self.smoke_expand_upward_diagonal_max_gap_ratio, ) if gap_x > max_gap: return False cx = (float(probe_box[0]) + float(probe_box[2])) / 2.0 return float(crop_x1) <= cx <= float(crop_x2) def _smoke_corroborates_expansion_seed( self, probe_box: np.ndarray, seed: BoundingBox, seed_arr: np.ndarray, crop_x1: int, crop_x2: int, ) -> bool: if self._boxes_involve_each_other_arrays( probe_box, seed_arr, self.smoke_expand_corroborate_involve_thresh, ): return True if self._smoke_corroborates_upward_diagonal_plume( probe_box, seed_arr, crop_x1, crop_x2 ): return True smoke_h = max(1.0, float(seed.y2 - seed.y1)) band_pad = smoke_h * self.smoke_expand_vertical_band_pad_ratio cy = (float(probe_box[1]) + float(probe_box[3])) / 2.0 cx = (float(probe_box[0]) + float(probe_box[2])) / 2.0 if cy < float(seed.y1) - band_pad or cy > float(seed.y2) + band_pad: return False gap_x = self._box_axis_gap(probe_box, seed_arr, 0) seed_w = max(1.0, float(seed.x2 - seed.x1)) max_band_gap = max( 32.0, seed_w * self.smoke_expand_band_max_horizontal_gap_ratio, ) if gap_x > max_band_gap: return False return float(crop_x1) <= cx <= float(crop_x2) def _smoke_box_border_strip(self, bw: int, bh: int) -> int: return max( 2, min( int(min(bw, bh) * self.smoke_expand_border_strip_ratio), bw // 3, bh // 3, ), ) def _smoke_expansion_side_border_rgb( self, image: np.ndarray, x1: int, y1: int, x2: int, y2: int, side: str, ) -> np.ndarray | None: bw = x2 - x1 bh = y2 - y1 if bw < 2 or bh < 2: return None strip = self._smoke_box_border_strip(bw, bh) if side == "left": block = image[y1:y2, x1 : x1 + strip] elif side == "right": block = image[y1:y2, x2 - strip : x2] elif side == "top": block = image[y1 : y1 + strip, x1:x2] else: block = image[y2 - strip : y2, x1:x2] if block.size == 0: return None return np.array( [ float(np.mean(block[:, :, 2])), float(np.mean(block[:, :, 1])), float(np.mean(block[:, :, 0])), ], dtype=np.float32, ) def _smoke_expansion_border_reference_rgb( self, image: np.ndarray, seed: BoundingBox ) -> np.ndarray | None: h, w = image.shape[:2] x1 = max(0, int(math.floor(seed.x1))) y1 = max(0, int(math.floor(seed.y1))) x2 = min(w, int(math.ceil(seed.x2))) y2 = min(h, int(math.ceil(seed.y2))) block_means: list[np.ndarray] = [] for side in ("top", "bottom", "left", "right"): rgb = self._smoke_expansion_side_border_rgb(image, x1, y1, x2, y2, side) if rgb is not None: block_means.append(rgb) if not block_means: return None return np.mean(block_means, axis=0) def _smoke_expansion_tta_view_corroborated( self, index: int, boxes: np.ndarray, cls_ids: np.ndarray, view_ids: np.ndarray, cls_smoke: int, ) -> bool: n = len(boxes) if n <= 1 or int(cls_ids[index]) != cls_smoke: return False iou = self._compute_iou_matrix(boxes) other = np.arange(n) != index other_view = view_ids != view_ids[index] same_class = cls_ids == cls_smoke return bool( np.any( other & other_view & same_class & (iou[index] >= self.smoke_expand_tta_view_iou_thresh) ) ) def _probe_extends_smoke_upward( self, probe_box: np.ndarray, seed_arr: np.ndarray, margin: float = 2.0 ) -> bool: return float(probe_box[1]) < float(seed_arr[1]) - margin def _probe_extends_smoke_downward( self, probe_box: np.ndarray, seed_arr: np.ndarray, margin: float = 2.0 ) -> bool: return float(probe_box[3]) > float(seed_arr[3]) + margin def _probe_qualifies_for_upward_low_conf( self, probe_box: np.ndarray, seed_arr: np.ndarray, margin: float = 2.0 ) -> bool: if not self._probe_extends_smoke_upward(probe_box, seed_arr, margin): return False return not self._probe_extends_smoke_downward(probe_box, seed_arr, margin) def _probe_extends_smoke_horizontally( self, probe_box: np.ndarray, seed_arr: np.ndarray, margin: float = 2.0 ) -> bool: return ( float(probe_box[0]) < float(seed_arr[0]) - margin or float(probe_box[2]) > float(seed_arr[2]) + margin ) def _probe_extends_smoke_seed( self, probe_box: np.ndarray, seed_arr: np.ndarray, margin: float = 2.0 ) -> bool: return ( self._probe_extends_smoke_horizontally(probe_box, seed_arr, margin) or self._probe_extends_smoke_upward(probe_box, seed_arr, margin) or float(probe_box[3]) > float(seed_arr[3]) + margin ) def _smoke_expansion_tta_candidate_conf_ok( self, score: float, candidate_box: np.ndarray, seed_arr: np.ndarray ) -> bool: if self._probe_qualifies_for_upward_low_conf(candidate_box, seed_arr): if score >= self.smoke_expand_min_upward_tta_conf: return True if score >= self.smoke_expand_min_upward_tta_involve_conf: return self._boxes_involve_each_other_arrays( candidate_box, seed_arr, self.smoke_expand_corroborate_involve_thresh, ) return False if score < self.smoke_expand_min_tta_smoke_conf: return False if self._probe_extends_smoke_horizontally( candidate_box, seed_arr ) or self._probe_extends_smoke_downward(candidate_box, seed_arr): return score >= self.smoke_expand_min_tta_extend_conf return True def _smoke_expansion_crop_candidate_conf_ok( self, score: float, candidate_box: np.ndarray, seed_arr: np.ndarray ) -> bool: if ( self._probe_extends_smoke_horizontally(candidate_box, seed_arr) and self._probe_extends_smoke_downward(candidate_box, seed_arr) and score < self.smoke_expand_min_bidi_crop_conf ): return False if self._probe_qualifies_for_upward_low_conf(candidate_box, seed_arr): return score >= self.smoke_expand_min_upward_crop_conf if score < self.smoke_expand_min_probe_smoke_conf: return False if self._probe_extends_smoke_horizontally( candidate_box, seed_arr ) or self._probe_extends_smoke_downward(candidate_box, seed_arr): return score >= self.smoke_expand_min_extend_probe_conf return True def _smoke_expand_mean_rgb_distance( self, image: np.ndarray, box: np.ndarray, ref_rgb: np.ndarray ) -> float: h, w = image.shape[:2] x1 = max(0, int(math.floor(float(box[0])))) y1 = max(0, int(math.floor(float(box[1])))) x2 = min(w, int(math.ceil(float(box[2])))) y2 = min(h, int(math.ceil(float(box[3])))) if x2 <= x1 or y2 <= y1: return float("inf") roi = image[y1:y2, x1:x2] if roi.size == 0: return float("inf") mean_rgb = np.array( [ float(np.mean(roi[:, :, 2])), float(np.mean(roi[:, :, 1])), float(np.mean(roi[:, :, 0])), ], dtype=np.float32, ) return float(np.linalg.norm(mean_rgb - ref_rgb)) def _passes_smoke_expand_probe_color( self, image: np.ndarray, probe_box: np.ndarray, ref_rgb: np.ndarray, ) -> bool: return ( self._smoke_expand_mean_rgb_distance(image, probe_box, ref_rgb) <= self.smoke_expand_max_color_dist ) def _smoke_expansion_union_candidate( self, union_box: np.ndarray, seed_arr: np.ndarray, candidate_box: np.ndarray, candidate_score: float, max_conf: float, ) -> tuple[np.ndarray, float, bool]: union_box = union_box.copy() union_box[0] = min(union_box[0], candidate_box[0]) union_box[1] = min(union_box[1], candidate_box[1]) union_box[2] = max(union_box[2], candidate_box[2]) union_box[3] = max(union_box[3], candidate_box[3]) margin = 2.0 found_extra = ( float(candidate_box[0]) < float(seed_arr[0]) - margin or float(candidate_box[2]) > float(seed_arr[2]) + margin or float(candidate_box[1]) < float(seed_arr[1]) - margin or float(candidate_box[3]) > float(seed_arr[3]) + margin ) return union_box, max(max_conf, candidate_score), found_extra def _collect_smoke_expansion_cluster( self, image: np.ndarray, seed: BoundingBox, orig_size: tuple[int, int], crop_x1: int, crop_x2: int, ref_rgb: np.ndarray | None, tta_boxes: np.ndarray, tta_scores: np.ndarray, tta_cls: np.ndarray, tta_view_ids: np.ndarray | None, crop_boxes: np.ndarray | None = None, crop_scores: np.ndarray | None = None, crop_cls: np.ndarray | None = None, ) -> tuple[np.ndarray, float] | None: cls_smoke = self.class_names.index("smoke") seed_arr = self._bbox_to_array(seed) union_box = seed_arr.copy() max_conf = float(seed.conf) found_extra = False for i in range(len(tta_boxes)): if int(tta_cls[i]) != cls_smoke: continue candidate_box = tta_boxes[i] if not self._smoke_expansion_tta_candidate_conf_ok( float(tta_scores[i]), candidate_box, seed_arr ): continue if tta_view_ids is not None and not self._smoke_expansion_tta_view_corroborated( i, tta_boxes, tta_cls, tta_view_ids, cls_smoke ): continue if not self._smoke_corroborates_expansion_seed( candidate_box, seed, seed_arr, crop_x1, crop_x2 ): continue if not self._smoke_probe_within_expansion_reach( seed_arr, candidate_box, orig_size ): continue union_box, max_conf, extra = self._smoke_expansion_union_candidate( union_box, seed_arr, candidate_box, float(tta_scores[i]), max_conf ) found_extra = found_extra or extra if ( crop_boxes is not None and crop_scores is not None and crop_cls is not None and ref_rgb is not None and len(crop_boxes) > 0 ): for i in range(len(crop_boxes)): if int(crop_cls[i]) != cls_smoke: continue candidate_box = crop_boxes[i] if not self._smoke_expansion_crop_candidate_conf_ok( float(crop_scores[i]), candidate_box, seed_arr ): continue if not self._smoke_corroborates_expansion_seed( candidate_box, seed, seed_arr, crop_x1, crop_x2 ): continue if not self._smoke_probe_within_expansion_reach( seed_arr, candidate_box, orig_size ): continue if not self._passes_smoke_expand_probe_color( image, candidate_box, ref_rgb ): continue candidate_box = self._clip_smoke_crop_candidate_for_union( candidate_box, seed_arr ) union_box, max_conf, extra = self._smoke_expansion_union_candidate( union_box, seed_arr, candidate_box, float(crop_scores[i]), max_conf, ) found_extra = found_extra or extra union_box = self._clamp_smoke_expansion_union(union_box, seed_arr, orig_size) margin = 2.0 refined_seed = ( float(union_box[0]) < float(seed_arr[0]) - margin or float(union_box[2]) > float(seed_arr[2]) + margin or float(union_box[1]) < float(seed_arr[1]) - margin or float(union_box[3]) > float(seed_arr[3]) + margin or float(union_box[0]) > float(seed_arr[0]) + margin or float(union_box[2]) < float(seed_arr[2]) - margin or float(union_box[1]) > float(seed_arr[1]) + margin or float(union_box[3]) < float(seed_arr[3]) - margin ) if not found_extra and not refined_seed: return None return union_box, max_conf def _smoke_expansion_seed_is_complete( self, seed: BoundingBox, orig_size: tuple[int, int], tta_boxes: np.ndarray, tta_cls: np.ndarray, tta_view_ids: np.ndarray | None, ) -> bool: orig_w, orig_h = orig_size seed_w = float(seed.x2 - seed.x1) width_ratio = seed_w / max(1.0, float(orig_w)) y1_ratio = float(seed.y1) / max(1.0, float(orig_h)) if y1_ratio > self.smoke_expand_skip_max_y1_ratio: return False if width_ratio < self.smoke_expand_skip_min_width_ratio: return False if tta_view_ids is None or len(tta_boxes) == 0: return False cls_smoke = self.class_names.index("smoke") seed_arr = self._bbox_to_array(seed) matched_views: set[int] = set() for i in range(len(tta_boxes)): if int(tta_cls[i]) != cls_smoke: continue if self._boxes_involve_each_other_arrays( tta_boxes[i], seed_arr, self.smoke_expand_corroborate_involve_thresh, ): matched_views.add(int(tta_view_ids[i])) if len(matched_views) < 2: return False return seed.conf >= self.smoke_expand_skip_min_conf @staticmethod def _smoke_expansion_union_extended_seed( seed_arr: np.ndarray, union_box: np.ndarray, margin: float = 2.0 ) -> bool: return ( float(union_box[0]) < float(seed_arr[0]) - margin or float(union_box[2]) > float(seed_arr[2]) + margin or float(union_box[1]) < float(seed_arr[1]) - margin or float(union_box[3]) > float(seed_arr[3]) + margin ) def _probe_smoke_expansion( self, image: np.ndarray, results: list[BoundingBox], orig_size: tuple[int, int], tta_boxes: np.ndarray, tta_scores: np.ndarray, tta_cls: np.ndarray, tta_view_ids: np.ndarray | None, ) -> list[BoundingBox]: if not self.use_tta: return results cls_smoke = self.class_names.index("smoke") seeds = sorted( ( box for box in results if box.cls_id == cls_smoke and box.conf >= self.smoke_expand_min_seed_conf ), key=lambda box: box.conf, reverse=True, ) if not seeds: return results seeds = seeds[: self.smoke_expand_max_seeds] absorbed: set[int] = set() replacements: dict[int, BoundingBox] = {} for seed in seeds: seed_key = id(seed) if seed_key in absorbed: continue skip_crops_only = self._smoke_expansion_seed_is_complete( seed, orig_size, tta_boxes, tta_cls, tta_view_ids ) crop_regions = self._smoke_expansion_crop_regions(seed, orig_size) primary_x1 = crop_regions[0][0] primary_x2 = crop_regions[0][1] ref_rgb = self._smoke_expansion_border_reference_rgb(image, seed) expanded = self._collect_smoke_expansion_cluster( image, seed, orig_size, primary_x1, primary_x2, ref_rgb, tta_boxes, tta_scores, tta_cls, tta_view_ids, None, None, None, ) seed_arr = self._bbox_to_array(seed) orig_h = orig_size[1] upward_open = float(seed.y1) / max(1.0, float(orig_h)) > ( self.smoke_expand_skip_max_y1_ratio ) if skip_crops_only: pass elif ( expanded is not None and self._smoke_expansion_union_extended_seed(seed_arr, expanded[0]) and not upward_open ): pass else: crop_boxes_list: list[np.ndarray] = [] crop_scores_list: list[np.ndarray] = [] crop_cls_list: list[np.ndarray] = [] for x1, x2, y1, y2 in crop_regions: crop = image[y1:y2, x1:x2] if crop.size == 0 or crop.shape[0] < 2 or crop.shape[1] < 2: continue crop_size = (crop.shape[1], crop.shape[0]) boxes, scores, cls_ids = self._infer_view_arrays( crop, crop_size, flip_code=None ) if len(boxes) == 0: continue boxes = boxes.copy() boxes[:, [0, 2]] += x1 boxes[:, [1, 3]] += y1 boxes = self._clip_boxes(boxes, orig_size) crop_boxes_list.append(boxes) crop_scores_list.append(scores) crop_cls_list.append(cls_ids) merged_crop_boxes: np.ndarray | None = None merged_crop_scores: np.ndarray | None = None merged_crop_cls: np.ndarray | None = None if crop_boxes_list and ref_rgb is not None: merged_crop_boxes = np.concatenate(crop_boxes_list, axis=0) merged_crop_scores = np.concatenate(crop_scores_list, axis=0) merged_crop_cls = np.concatenate(crop_cls_list, axis=0) crop_expanded = self._collect_smoke_expansion_cluster( image, seed, orig_size, primary_x1, primary_x2, ref_rgb, tta_boxes, tta_scores, tta_cls, tta_view_ids, merged_crop_boxes, merged_crop_scores, merged_crop_cls, ) if crop_expanded is not None: expanded = crop_expanded if expanded is None: continue union_box, max_conf = expanded seed_arr = self._bbox_to_array(seed) for other in results: if other.cls_id != cls_smoke or other is seed: continue other_key = id(other) if other_key in absorbed: continue if other_key in replacements: other_arr = self._bbox_to_array(replacements[other_key]) else: other_arr = self._bbox_to_array(other) if self._boxes_involve_each_other_arrays( other_arr, seed_arr, self.smoke_expand_corroborate_involve_thresh, ) or self._boxes_involve_each_other_arrays( other_arr, union_box, self.smoke_expand_corroborate_involve_thresh, ): union_box[0] = min(union_box[0], other_arr[0]) union_box[1] = min(union_box[1], other_arr[1]) union_box[2] = max(union_box[2], other_arr[2]) union_box[3] = max(union_box[3], other_arr[3]) other_conf = ( float(replacements[other_key].conf) if other_key in replacements else float(other.conf) ) max_conf = max(max_conf, other_conf) absorbed.add(other_key) expanded_smoke = BoundingBox( x1=int(math.floor(union_box[0])), y1=int(math.floor(union_box[1])), x2=int(math.ceil(union_box[2])), y2=int(math.ceil(union_box[3])), cls_id=cls_smoke, conf=min(1.0, max_conf), ) replacements[seed_key] = expanded_smoke if not replacements and not absorbed: return results updated: list[BoundingBox] = [] for box in results: box_key = id(box) if box_key in absorbed: continue if box_key in replacements: updated.append(replacements[box_key]) continue updated.append(box) return updated def _infer_view_arrays( self, image: np.ndarray, orig_size: tuple[int, int], flip_code: int | None = None, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: if flip_code is not None: image = cv2.flip(image, flip_code) decode_size = (image.shape[1], image.shape[0]) input_tensor, ratio, pad, _ = self._preprocess(image) expected = (1, 3, self.input_height, self.input_width) if input_tensor.shape != expected: raise ValueError( f"Bad input tensor shape={input_tensor.shape}, expected={expected}" ) outputs = self.session.run(self.output_names, {self.input_name: input_tensor}) boxes, scores, cls_ids = self._decode_preds_to_arrays( outputs[0], ratio, pad, decode_size ) if flip_code is not None and len(boxes) > 0: boxes = self._map_flipped_boxes_arrays(boxes, decode_size, flip_code) return boxes, scores, cls_ids def _collect_merged_arrays( self, image: np.ndarray ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, tuple[int, int]]: orig_size = (image.shape[1], image.shape[0]) if self.use_tta: view_flip_codes: tuple[int | None, ...] = (None, 1) else: view_flip_codes = (None,) all_boxes: list[np.ndarray] = [] all_scores: list[np.ndarray] = [] all_cls: list[np.ndarray] = [] all_view_ids: list[np.ndarray] = [] view_id = 0 for flip_code in view_flip_codes: boxes, scores, cls_ids = self._infer_view_arrays( image, orig_size, flip_code ) if len(boxes) > 0: all_boxes.append(boxes) all_scores.append(scores) all_cls.append(cls_ids) all_view_ids.append( np.full(len(boxes), view_id, dtype=np.int32) ) view_id += 1 if not all_boxes: return ( np.empty((0, 4), dtype=np.float32), np.empty((0,), dtype=np.float32), np.empty((0,), dtype=np.int32), np.empty((0,), dtype=np.int32), orig_size, ) return ( np.concatenate(all_boxes, axis=0), np.concatenate(all_scores, axis=0), np.concatenate(all_cls, axis=0), np.concatenate(all_view_ids, axis=0), orig_size, ) def _decode_final_dets( self, preds: np.ndarray, ratio: float, pad: tuple[float, float], orig_size: tuple[int, int], ) -> list[BoundingBox]: boxes, scores, cls_ids = self._decode_preds_to_arrays( preds, ratio, pad, orig_size ) return self._apply_post_filters(boxes, scores, cls_ids, orig_size) def _postprocess( self, output: np.ndarray, ratio: float, pad: tuple[float, float], orig_size: tuple[int, int], ) -> list[BoundingBox]: return self._decode_final_dets(output, ratio, pad, orig_size) @staticmethod def _validate_image(image: np.ndarray) -> np.ndarray: if image is None: raise ValueError("Input image is None") if not isinstance(image, np.ndarray): raise TypeError(f"Input is not numpy array: {type(image)}") if image.ndim != 3: raise ValueError(f"Expected HWC image, got shape={image.shape}") if image.shape[0] <= 0 or image.shape[1] <= 0: raise ValueError(f"Invalid image shape={image.shape}") if image.shape[2] != 3: raise ValueError(f"Expected 3 channels, got shape={image.shape}") if image.dtype != np.uint8: image = image.astype(np.uint8) return image def _predict_single(self, image: np.ndarray) -> list[BoundingBox]: image = self._validate_image(image) boxes, scores, cls_ids, view_ids, orig_size = self._collect_merged_arrays( image ) results = self._apply_post_filters( boxes, scores, cls_ids, orig_size, view_ids, ) cls_smoke = self.class_names.index("smoke") has_expandable_smoke = any( box.cls_id == cls_smoke and box.conf >= self.smoke_expand_min_seed_conf for box in results ) smoke_seeds = [ BoundingBox( x1=box.x1, y1=box.y1, x2=box.x2, y2=box.y2, cls_id=box.cls_id, conf=box.conf, ) for box in results if box.cls_id == cls_smoke ] if has_expandable_smoke: results = self._probe_smoke_expansion( image, results, orig_size, boxes, scores, cls_ids, view_ids ) has_anchor_smoke = any( box.cls_id == cls_smoke and box.conf >= self.smoke_anchor_min_smoke_conf for box in results ) if has_anchor_smoke: results = self._probe_smoke_anchored_fire( image, results, orig_size, boxes, cls_ids, smoke_seeds=smoke_seeds ) results = self._probe_merged_anchored_fire_ext( image, results, orig_size, boxes, scores, cls_ids ) results = self._filter_probe_fires_by_color(image, results) results = self._filter_fire_ext_by_red_color(image, results) return self._filter_results_min_conf(results) def predict_batch( self, batch_images: list[ndarray], offset: int, n_keypoints: int, ) -> list[TVFrameResult]: results: list[TVFrameResult] = [] for frame_number_in_batch, image in enumerate(batch_images): try: boxes = self._predict_single(image) except Exception as e: print( f"⚠️ Inference failed for frame " f"{offset + frame_number_in_batch}: {e}" ) boxes = [] results.append( TVFrameResult( frame_id=offset + frame_number_in_batch, boxes=boxes, keypoints=[(0, 0) for _ in range(max(0, int(n_keypoints)))], ) ) return results