""" MicroGhost-Thermal: Inference & Deployment Module (V2) ======================================================== Handles edge inference, multi-target detection, NMS, alert generation, and model export (ONNX, TFLite/INT8) for ESP32-S3 deployment. V2 Additions: - Supports loading MicroGhostV2 with BiFusion neck and EnergyGate - Dual threshold logic (inference vs eval) - Gate weight extraction during inference - Graceful backward compatibility for V1 models """ import os import time import json import numpy as np import torch import torch.nn as nn import cv2 from model import MicroGhostThermal, MicroGhostV2, estimate_peak_sram from preprocessing import ThermalPreprocessor from config import ( INPUT_SIZE, INPUT_WIDTH, INPUT_HEIGHT, CLASS_MAP, NUM_CLASSES, NUM_ANCHORS, DEFAULT_ANCHOR_SIZES, CONFIDENCE_THRESHOLD, EVAL_CONFIDENCE_THRESHOLD, NMS_IOU_THRESHOLD, MAX_DETECTIONS, DEVICE, ESP32_S3, ALERT_CONFIG, V2_CLASSIFIER_HIDDEN_DIM, CLASSIFIER_HIDDEN_DIM, MIN_BOX_WIDTH_NORM, MIN_BOX_HEIGHT_NORM, MIN_BOX_AREA_NORM, LOG_CLAMP_MIN, LOG_CLAMP_MAX ) # Inverse class map for alert generation INV_CLASS_MAP = {v: k for k, v in CLASS_MAP.items()} def _normalize_legacy_state_dict(state_dict): """Map checkpoints saved with the older GhostModule layout to this model.""" normalized = {} for key, value in state_dict.items(): new_key = key new_key = new_key.replace(".primary_conv.0.", ".primary_conv.") new_key = new_key.replace(".primary_conv.1.", ".primary_bn.") normalized[new_key] = value return normalized def _infer_checkpoint_num_anchors(state_dict): """Infer anchor count from the detection head tensors.""" obj_weight = state_dict.get("head_small.obj_head.weight") if obj_weight is not None: return int(obj_weight.shape[0]) bbox_weight = state_dict.get("head_small.bbox_head.weight") if bbox_weight is not None and bbox_weight.shape[0] % 4 == 0: return int(bbox_weight.shape[0] // 4) return None def _checkpoint_is_v2(config, state_dict): return ( config.get("model_version", "v1") == "v2" or any(k.startswith(("energy_gate.", "bifusion_neck.", "thm_stem.")) for k in state_dict) ) def _match_classifier_to_checkpoint(model, state_dict): """Resize the final classifier layer when loading legacy classifier heads.""" weight = state_dict.get("classifier.classifier.3.weight") bias = state_dict.get("classifier.classifier.3.bias") if weight is None or bias is None: return final_layer = model.classifier.classifier[3] checkpoint_outputs = int(weight.shape[0]) if checkpoint_outputs != final_layer.out_features: model.classifier.classifier[3] = nn.Linear(final_layer.in_features, checkpoint_outputs) # ============================================================================ # 1. MODEL LOADING # ============================================================================ def load_inference_model(model_path, device=DEVICE, override_num_anchors=None): """Load PyTorch model for inference, supporting both V1 and V2 architectures.""" if not os.path.exists(model_path): raise FileNotFoundError(f"Model file not found: {model_path}") print(f"Loading model from {model_path}...") checkpoint = torch.load(model_path, map_location=device, weights_only=False) config = checkpoint.get('config', {}) state_dict = checkpoint.get('model_state_dict', checkpoint) state_dict = _normalize_legacy_state_dict(state_dict) # Determine model version is_v2 = _checkpoint_is_v2(config, state_dict) checkpoint_anchors = _infer_checkpoint_num_anchors(state_dict) anchors_to_use = config.get('num_anchors') or checkpoint_anchors or NUM_ANCHORS if override_num_anchors is not None: if checkpoint_anchors is not None and checkpoint_anchors != override_num_anchors: print( f"Warning: requested {override_num_anchors} anchors, " f"but checkpoint uses {checkpoint_anchors}; using checkpoint value." ) anchors_to_use = checkpoint_anchors else: anchors_to_use = override_num_anchors if is_v2: model = MicroGhostV2( num_classes=config.get('num_classes', NUM_CLASSES), input_size=config.get('input_size', INPUT_SIZE), classifier_hidden_dim=config.get('classifier_hidden_dim', V2_CLASSIFIER_HIDDEN_DIM), num_anchors=anchors_to_use, training_mode=False ) else: model = MicroGhostThermal( num_classes=config.get('num_classes', NUM_CLASSES), input_size=config.get('input_size', INPUT_SIZE), classifier_hidden_dim=config.get('classifier_hidden_dim', CLASSIFIER_HIDDEN_DIM), num_anchors=anchors_to_use ) _match_classifier_to_checkpoint(model, state_dict) missing, unexpected = model.load_state_dict(state_dict, strict=False) model.to(device) model.eval() print(f"Loaded {'MicroGhost-V2' if is_v2 else 'MicroGhost-V1'} successfully.") if missing or unexpected: print( "Compatibility load note: " f"{len(missing)} missing/new keys and {len(unexpected)} unused keys." ) return model, anchors_to_use # ============================================================================ # 2. INFERENCE ENGINE (MULTI-TARGET + NMS) # ============================================================================ def calculate_iou_numpy(box1, box2): """Calculate IoU between two bounding boxes [x1, y1, x2, y2].""" x1 = max(box1[0], box2[0]) y1 = max(box1[1], box2[1]) x2 = min(box1[2], box2[2]) y2 = min(box1[3], box2[3]) inter_area = max(0, x2 - x1) * max(0, y2 - y1) if inter_area == 0: return 0.0 box1_area = (box1[2] - box1[0]) * (box1[3] - box1[1]) box2_area = (box2[2] - box2[0]) * (box2[3] - box2[1]) return inter_area / (box1_area + box2_area - inter_area) def nms(detections, iou_threshold=NMS_IOU_THRESHOLD): """Non-Maximum Suppression to filter overlapping boxes.""" if not detections: return [] # Sort by confidence descending detections = sorted(detections, key=lambda x: x['conf'], reverse=True) keep = [] for det in detections: discard = False for kept_det in keep: iou = calculate_iou_numpy(det['bbox'], kept_det['bbox']) if iou > iou_threshold: discard = True break if not discard: keep.append(det) return keep[:MAX_DETECTIONS] def filter_spurious_detections(detections): """Remove tiny boxes and edge artifacts.""" kept = [] for det in detections: x1, y1, x2, y2 = det['bbox'] bw = x2 - x1 bh = y2 - y1 area = bw * bh if bw < MIN_BOX_WIDTH_NORM or bh < MIN_BOX_HEIGHT_NORM: continue if area < MIN_BOX_AREA_NORM: continue kept.append(det) return kept def decode_predictions(obj_small, bbox_small, obj_large, bbox_large, anchor_sizes=None, is_eval=False, iou_pred=1.0, conf_threshold=None): if anchor_sizes is None: anchor_sizes = DEFAULT_ANCHOR_SIZES if getattr(__import__('config'), 'DEBUG_MODE', False): conf_threshold = 0.05 else: conf_threshold = ( conf_threshold if conf_threshold is not None else EVAL_CONFIDENCE_THRESHOLD if is_eval else CONFIDENCE_THRESHOLD ) small_grid_w, small_grid_h = INPUT_WIDTH // 8, INPUT_HEIGHT // 8 large_grid_w, large_grid_h = INPUT_WIDTH // 16, INPUT_HEIGHT // 16 def _decode_box(pred_box, grid_x, grid_y, grid_w, grid_h, anchor_size): cx = (grid_x + pred_box[0]) / grid_w cy = (grid_y + pred_box[1]) / grid_h w = anchor_size * np.exp(np.clip(pred_box[2], LOG_CLAMP_MIN, LOG_CLAMP_MAX)) h = anchor_size * np.exp(np.clip(pred_box[3], LOG_CLAMP_MIN, LOG_CLAMP_MAX)) return [ max(0.0, min(1.0, cx - w / 2)), max(0.0, min(1.0, cy - h / 2)), max(0.0, min(1.0, cx + w / 2)), max(0.0, min(1.0, cy + h / 2)), ] def _extract(obj_map, bbox_map, grid_w, grid_h): candidates = [] obj_probs = torch.sigmoid(obj_map).cpu().numpy() bbox_data = bbox_map.cpu().numpy() for a in range(obj_probs.shape[0]): y_idx, x_idx = np.where(obj_probs[a] > conf_threshold) for gy, gx in zip(y_idx, x_idx): off = a * 4 bbox = _decode_box( bbox_data[off:off + 4, gy, gx], gx, gy, grid_w, grid_h, anchor_sizes[a], ) candidates.append({ 'conf': float(obj_probs[a, gy, gx]) * iou_pred, 'bbox': bbox, }) return candidates candidates = [] candidates.extend(_extract(obj_small, bbox_small, small_grid_w, small_grid_h)) candidates.extend(_extract(obj_large, bbox_large, large_grid_w, large_grid_h)) candidates = filter_spurious_detections(candidates) return nms(candidates) class ThermalInferenceEngine: def __init__(self, model_path=None, model=None, device=DEVICE, override_num_anchors=None): self.device = device self.preprocessor = ThermalPreprocessor() if model is not None: self.model = model.to(self.device) self.model.eval() self.anchor_sizes = DEFAULT_ANCHOR_SIZES elif model_path is not None: self.model, anchors_used = load_inference_model(model_path, device, override_num_anchors) if anchors_used == 2: self.anchor_sizes = [0.108, 0.180] # Fallback V1 sizes else: self.anchor_sizes = DEFAULT_ANCHOR_SIZES else: raise ValueError("Must provide either model_path or model instance.") def detect(self, image_rgb, image_thermal, is_eval=False, conf_threshold=None): """ Run full detection pipeline on a dual-modality input. Args: image_rgb: numpy array (H, W, 3) image_thermal: numpy array (H, W) is_eval: If true, uses the lower EVAL_CONFIDENCE_THRESHOLD to maximize recall. Returns: list of dicts: {'conf': float, 'bbox': [x1,y1,x2,y2], 'class': str, 'gate_w': dict} """ h_orig, w_orig = image_thermal.shape[:2] # 1. Preprocess img_tensor, _ = self.preprocessor.process( image_rgb, image_thermal, [], [], img_size=(h_orig, w_orig), augment=False ) img_tensor = img_tensor.unsqueeze(0).to(self.device) # 2. Forward pass with torch.no_grad(): preds = self.model(img_tensor) # 3. Classifier predictions label_outputs = preds['label'].shape[1] if label_outputs == NUM_CLASSES + 1: cls_logits = preds['label'][:, :NUM_CLASSES] iou_pred = float(torch.sigmoid(preds['label'][:, NUM_CLASSES])[0].item()) else: cls_logits = preds['label'] iou_pred = 1.0 cls_probs = torch.softmax(cls_logits, dim=1)[0].cpu().numpy() pred_class_idx = int(np.argmax(cls_probs)) # 4. Extract boxes from both scales with IoU-aware confidence final_detections = decode_predictions( preds['obj_small'][0], preds['bbox_small'][0], preds['obj_large'][0], preds['bbox_large'][0], anchor_sizes=self.anchor_sizes, is_eval=is_eval, iou_pred=iou_pred, conf_threshold=conf_threshold, ) # If classifier says background, return empty (unless in strict eval where we might bypass) if pred_class_idx == CLASS_MAP['background']: return [] detected_class_name = INV_CLASS_MAP.get(pred_class_idx, "unknown") class_conf = float(cls_probs[pred_class_idx]) # Extract mean gate weights for diagnostic (V2 only) gate_w = None if 'w_rgb' in preds: gate_w = { 'rgb': float(preds['w_rgb'][0].mean().item()), 'thm': float(preds['w_thm'][0].mean().item()) } for det in final_detections: det['bbox'] = [float(v) for v in det['bbox']] det['class'] = detected_class_name det['combined_conf'] = (det['conf'] + class_conf) / 2.0 det['gate_w'] = gate_w # Pseudo-radiometric temp extraction x1, y1, x2, y2 = det['bbox'] x1_p, y1_p = int(x1 * w_orig), int(y1 * h_orig) x2_p, y2_p = int(x2 * w_orig), int(y2 * h_orig) x1_p, y1_p = max(0, x1_p), max(0, y1_p) x2_p, y2_p = min(w_orig, x2_p), min(h_orig, y2_p) if x2_p > x1_p and y2_p > y1_p: crop = image_thermal[y1_p:y2_p, x1_p:x2_p] max_val = np.max(crop) det['temp_c'] = round(20.0 + (max_val / 255.0) * 20.0, 1) else: det['temp_c'] = 0.0 return final_detections def filter_by_laplacian(self, lap_image, detections, lap_thresh=80.0, high_conf_bypass=None): """Reject detections whose visual crop is too smooth to look like a real target.""" if lap_thresh is None: return detections h_lap, w_lap = lap_image.shape[:2] confirmed = [] for det in detections: x1, y1, x2, y2 = det['bbox'] x1_p = max(0, min(w_lap, int(x1 * w_lap))) y1_p = max(0, min(h_lap, int(y1 * h_lap))) x2_p = max(0, min(w_lap, int(x2 * w_lap))) y2_p = max(0, min(h_lap, int(y2 * h_lap))) det = dict(det) if x2_p <= x1_p or y2_p <= y1_p: det['lap_var'] = 0.0 continue crop = lap_image[y1_p:y2_p, x1_p:x2_p] if crop.ndim == 3: gray = cv2.cvtColor(crop, cv2.COLOR_RGB2GRAY) else: gray = crop lap_var = float(cv2.Laplacian(gray, cv2.CV_64F).var()) det['lap_var'] = round(lap_var, 2) if high_conf_bypass is not None and det.get('combined_conf', 0.0) >= high_conf_bypass: det['lap_filter'] = 'bypassed_high_conf' confirmed.append(det) elif lap_var >= lap_thresh: det['lap_filter'] = 'kept' confirmed.append(det) return confirmed def merge_related_detections(self, detections, x_overlap_thresh=0.35, vertical_gap_thresh=0.08): """Merge boxes that look like stacked/overlapping parts of the same person.""" if len(detections) < 2: return detections def should_merge(a, b): ax1, ay1, ax2, ay2 = a['bbox'] bx1, by1, bx2, by2 = b['bbox'] aw, bw = ax2 - ax1, bx2 - bx1 ah, bh = ay2 - ay1, by2 - by1 if aw <= 0 or bw <= 0 or ah <= 0 or bh <= 0: return False x_overlap = max(0.0, min(ax2, bx2) - max(ax1, bx1)) x_overlap_ratio = x_overlap / max(min(aw, bw), 1e-6) cx_gap = abs(((ax1 + ax2) / 2.0) - ((bx1 + bx2) / 2.0)) same_column = x_overlap_ratio >= x_overlap_thresh or cx_gap <= min(aw, bw) * 0.65 y_overlap = max(0.0, min(ay2, by2) - max(ay1, by1)) y_gap = max(0.0, max(ay1, by1) - min(ay2, by2)) vertically_related = y_overlap > 0.0 or y_gap <= vertical_gap_thresh return same_column and vertically_related parent = list(range(len(detections))) def find(i): while parent[i] != i: parent[i] = parent[parent[i]] i = parent[i] return i def union(i, j): ri, rj = find(i), find(j) if ri != rj: parent[rj] = ri for i in range(len(detections)): for j in range(i + 1, len(detections)): if should_merge(detections[i], detections[j]): union(i, j) groups = {} for i, det in enumerate(detections): groups.setdefault(find(i), []).append(det) merged = [] for group in groups.values(): if len(group) == 1: merged.append(group[0]) continue best = max(group, key=lambda d: d.get('combined_conf', d.get('conf', 0.0))) out = dict(best) out['bbox'] = [ min(d['bbox'][0] for d in group), min(d['bbox'][1] for d in group), max(d['bbox'][2] for d in group), max(d['bbox'][3] for d in group), ] out['conf'] = max(d.get('conf', 0.0) for d in group) out['combined_conf'] = max(d.get('combined_conf', 0.0) for d in group) out['temp_c'] = max(d.get('temp_c', 0.0) for d in group) out['lap_var'] = max(d.get('lap_var', 0.0) for d in group) out['merged_parts'] = len(group) merged.append(out) return sorted(merged, key=lambda d: d.get('combined_conf', d.get('conf', 0.0)), reverse=True) def filter_thermal_only_artifacts(self, detections): """Remove common thermal-only false positives such as top-edge hot bonnets.""" kept = [] for det in detections: x1, y1, x2, y2 = [float(v) for v in det['bbox']] w = x2 - x1 h = y2 - y1 conf = float(det.get('combined_conf', det.get('conf', 0.0))) temp_c = float(det.get('temp_c', 0.0)) touches_top_artifact_band = y1 <= 0.08 and y2 <= 0.30 wide_cool_low_conf_artifact = w >= 0.12 and conf < 0.70 and temp_c < 39.0 too_flat = h <= 0 or w <= 0 if touches_top_artifact_band or wide_cool_low_conf_artifact or too_flat: continue kept.append(det) return kept def detect_confirmed(self, image_rgb, image_thermal, lap_thresh=80.0, lap_image=None, high_conf_bypass=None, is_eval=False, conf_threshold=None): """Run detection, then reject flat crops such as hot car bonnets.""" detections = self.detect( image_rgb, image_thermal, is_eval=is_eval, conf_threshold=conf_threshold, ) detections = self.filter_by_laplacian( lap_image if lap_image is not None else image_rgb, detections, lap_thresh=lap_thresh, high_conf_bypass=high_conf_bypass, ) return self.merge_related_detections(detections) # ============================================================================ # 3. ALERT GENERATOR & EXPORT # ============================================================================ class AlertGenerator: """Generates payloads for transmission from edge node.""" def __init__(self, config=ALERT_CONFIG): self.config = config def generate_json_alert(self, detections, node_id="NODE-001"): if not detections: return None alert = {"node": node_id, "status": "INTRUSION_DETECTED"} if self.config['include_timestamp']: alert["timestamp"] = int(time.time()) if self.config['include_intrusion_count']: alert["count"] = len(detections) targets = [] for d in detections: t = {} if self.config['include_confidence']: t["conf"] = round(d['combined_conf'], 3) t["class"] = d.get('class', 'unknown') if self.config['include_bbox']: t["bbox"] = [round(x, 3) for x in d['bbox']] if self.config['include_temperature']: t["temp_c"] = d.get('temp_c', 0.0) if 'gate_w' in d and d['gate_w'] is not None: t['gate_w'] = d['gate_w'] targets.append(t) alert["targets"] = targets return json.dumps(alert) def export_to_onnx(model, save_path): """Export PyTorch model to ONNX.""" os.makedirs(os.path.dirname(save_path) or '.', exist_ok=True) model.eval() model.to('cpu') dummy_input = torch.randn(1, 4, INPUT_HEIGHT, INPUT_WIDTH) # V2 has w_rgb, w_thm outputs. Collect actual output keys from a forward pass with torch.no_grad(): out = model(dummy_input) output_names = list(out.keys()) torch.onnx.export( model, dummy_input, save_path, export_params=True, opset_version=11, do_constant_folding=True, input_names=['input'], output_names=output_names, dynamic_axes={'input': {0: 'batch_size'}} ) print(f"Exported ONNX model to {save_path}") return save_path def benchmark_model(model): """Run performance benchmarks on the model.""" print("\n" + "="*50) print("RUNNING PERFORMANCE BENCHMARK (V2)") print("="*50) model.eval() model.to('cpu') sram_info = estimate_peak_sram(model) print("Memory Estimates:") print(f" Flash (INT8 weights): ~{sram_info['peak_activation_int8_kb'] * 2:.1f} KB") print(f" Peak SRAM Arena: {sram_info['total_arena_int8_kb']:.1f} KB") print(f" Fits ESP32-S3: {'YES' if sram_info['fits_esp32_s3'] else 'NO'}") dummy = torch.randn(1, 4, INPUT_HEIGHT, INPUT_WIDTH) with torch.no_grad(): for _ in range(10): model(dummy) iters = 100 t0 = time.time() with torch.no_grad(): for _ in range(iters): model(dummy) t1 = time.time() ms_per_frame = ((t1 - t0) / iters) * 1000 print(f" Average Time: {ms_per_frame:.2f} ms") print(f" Throughput: {iters / (t1 - t0):.1f} FPS") print("="*50) if __name__ == '__main__': print("Inference Module (V2) Test:") gen = AlertGenerator() dummy_dets = [ {'conf': 0.85, 'combined_conf': 0.9, 'bbox': [0.1, 0.2, 0.3, 0.4], 'temp_c': 35.2, 'gate_w': {'rgb': 0.8, 'thm': 0.2}}, ] print(gen.generate_json_alert(dummy_dets))