""" MicroGhost-Thermal: Preprocessing Module (V2) =============================================== Handles data transformation, normalization, augmentation, and grid-based target encoding for thermal intrusion detection. V2 Changes: - CMM (Causal Mode Multiplexer) augmentation modes - Log clamp range updated to [-4.5, 4.5] - Support for single-modality inputs (RGB-only or thermal-only) """ import cv2 import math import torch import numpy as np from config import ( INPUT_SIZE, INPUT_CHANNELS, NUM_CLASSES, NUM_ANCHORS, DEFAULT_ANCHOR_RATIOS, DEFAULT_ANCHOR_SIZES, SMALL_GRID_W, SMALL_GRID_H, LARGE_GRID_W, LARGE_GRID_H, INPUT_WIDTH, INPUT_HEIGHT, LOG_CLAMP_MIN, LOG_CLAMP_MAX, ) # ============================================================================ # GRID ENCODER # ============================================================================ class GridEncoder: """ Grid-based target encoder for dual-head detection. Uses log-space width/height encoding clamped to [LOG_CLAMP_MIN, LOG_CLAMP_MAX]. """ def __init__(self, num_anchors=None, anchor_ratios=None, anchor_sizes=None): self.num_anchors = num_anchors or NUM_ANCHORS if anchor_ratios is not None: self.ratios = torch.tensor(anchor_ratios, dtype=torch.float32) else: self.ratios = torch.tensor(DEFAULT_ANCHOR_RATIOS, dtype=torch.float32) if anchor_sizes is not None: self.anchor_sizes = torch.tensor(anchor_sizes, dtype=torch.float32) else: self.anchor_sizes = torch.tensor(DEFAULT_ANCHOR_SIZES, dtype=torch.float32) # Soft margin thresholds for scale assignment self.small_threshold = 0.10 self.large_threshold = 0.20 def update_anchors(self, anchor_ratios, anchor_sizes): """Update anchors after K-Means analysis.""" self.ratios = torch.tensor(anchor_ratios, dtype=torch.float32) self.anchor_sizes = torch.tensor(anchor_sizes, dtype=torch.float32) print(f"GridEncoder anchors updated!") print(f" Ratios: {anchor_ratios}") print(f" Sizes: {anchor_sizes}") def encode(self, boxes, labels, input_size=None): """ Convert raw boxes to grid targets. Args: boxes: List of [cx, cy, w, h] (normalized 0-1) labels: List of class_ids input_size: Override input size (unused, grid sizes from config) Returns: dict with bbox_small, obj_small, bbox_large, obj_large, label """ targets = { 'bbox_small': torch.zeros(self.num_anchors * 4, SMALL_GRID_H, SMALL_GRID_W), 'obj_small': torch.zeros(self.num_anchors, SMALL_GRID_H, SMALL_GRID_W), 'bbox_large': torch.zeros(self.num_anchors * 4, LARGE_GRID_H, LARGE_GRID_W), 'obj_large': torch.zeros(self.num_anchors, LARGE_GRID_H, LARGE_GRID_W), 'label': torch.tensor(0).long(), } if len(boxes) == 0: return targets # Dominant target: largest bounding box box_areas = [b[2] * b[3] for b in boxes] dominant_idx = box_areas.index(max(box_areas)) targets['label'] = torch.tensor(labels[dominant_idx]).long() for box, label in zip(boxes, labels): cx, cy, w, h = box box_area = w * h # Soft margin scale assignment scale_assignments = [] if box_area < self.large_threshold: scale_assignments.append('small') if box_area >= self.small_threshold: scale_assignments.append('large') if not scale_assignments: scale_assignments = ['small'] for scale in scale_assignments: if scale == 'small': grid_w, grid_h = SMALL_GRID_W, SMALL_GRID_H key_bbox = 'bbox_small' key_obj = 'obj_small' else: grid_w, grid_h = LARGE_GRID_W, LARGE_GRID_H key_bbox = 'bbox_large' key_obj = 'obj_large' grid_x = min(int(cx * grid_w), grid_w - 1) grid_y = min(int(cy * grid_h), grid_h - 1) # Best anchor (match aspect ratio) box_ratio = h / (w + 1e-6) ratio_diffs = torch.abs(self.ratios - box_ratio) anchor_idx = torch.argmin(ratio_diffs).item() targets[key_obj][anchor_idx, grid_y, grid_x] = 1.0 anchor_size = self.anchor_sizes[anchor_idx].item() off = anchor_idx * 4 targets[key_bbox][off + 0, grid_y, grid_x] = \ (cx * grid_w) - grid_x targets[key_bbox][off + 1, grid_y, grid_x] = \ (cy * grid_h) - grid_y # V2: use configurable clamp range targets[key_bbox][off + 2, grid_y, grid_x] = \ torch.log(torch.tensor(w / anchor_size + 1e-6)).clamp( LOG_CLAMP_MIN, LOG_CLAMP_MAX) targets[key_bbox][off + 3, grid_y, grid_x] = \ torch.log(torch.tensor(h / anchor_size + 1e-6)).clamp( LOG_CLAMP_MIN, LOG_CLAMP_MAX) return targets # ============================================================================ # THERMAL-SPECIFIC AUGMENTATIONS (with CMM support) # ============================================================================ class ThermalAugmentor: """ Data augmentation pipeline with V2 CMM (Causal Mode Multiplexer) support. CMM Modes: - ROTO: Both modalities receive real data (normal) - RXTO: Thermal zeroed, RGB real (forces RGB-only detection) - ROTX: RGB zeroed, Thermal real (forces Thermal-only detection) The CMM mode is applied BEFORE spatial augmentations to ensure the model learns from each modality independently. """ def __init__(self): self.input_h = INPUT_HEIGHT self.input_w = INPUT_WIDTH self._try_load_albumentations() def _try_load_albumentations(self): """Try to load albumentations; fall back to manual augmentation.""" try: import albumentations as A self.A = A self.use_albumentations = True self.train_transform = A.Compose([ A.HorizontalFlip(p=0.5), A.Affine( scale=(0.9, 1.1), translate_percent={"x": (-0.05, 0.05), "y": (-0.05, 0.05)}, p=0.5 ), A.RandomBrightnessContrast( brightness_limit=0.3, contrast_limit=0.3, p=0.7, ), A.GaussNoise(p=0.5), A.GaussianBlur(blur_limit=(3, 5), p=0.2), A.Resize(self.input_h, self.input_w), ], additional_targets={'image_thermal': 'image'}, bbox_params=A.BboxParams( format='pascal_voc', label_fields=['labels'], min_visibility=0.3, )) self.val_transform = A.Compose([ A.Resize(self.input_h, self.input_w), ], additional_targets={'image_thermal': 'image'}, bbox_params=A.BboxParams( format='pascal_voc', label_fields=['labels'], )) except ImportError: self.use_albumentations = False print("albumentations not installed. Using basic augmentation.") def apply_cmm(self, image_rgb, image_thermal, cmm_mode='roto'): """ Apply Causal Mode Multiplexer. Args: image_rgb: (H, W, 3) numpy uint8 image_thermal: (H, W) numpy uint8 cmm_mode: 'roto' (both), 'rxto' (thermal zeroed), 'rotx' (RGB zeroed) Returns: image_rgb, image_thermal (potentially zeroed) """ if cmm_mode == 'rxto': image_thermal = np.zeros_like(image_thermal) elif cmm_mode == 'rotx': image_rgb = np.zeros_like(image_rgb) return image_rgb, image_thermal def augment_train(self, image_rgb, image_thermal, bboxes, labels, cmm_alpha=0.0): """ Training augmentation with optional CMM. Args: cmm_alpha: probability of single-modality training. If > 0, with probability cmm_alpha, one modality is zeroed. Split evenly between RXTO and ROTX. """ # CMM modality dropout (replaces the old fixed 15%/5% dropout) if cmm_alpha > 0: r = np.random.random() if r < cmm_alpha / 2: image_thermal = np.zeros_like(image_thermal) elif r < cmm_alpha: image_rgb = np.zeros_like(image_rgb) if self.use_albumentations and bboxes: try: result = self.train_transform( image=image_rgb, image_thermal=image_thermal, bboxes=bboxes, labels=labels ) return (result['image'], result['image_thermal'], result['bboxes'], result['labels']) except Exception: pass return self._manual_augment(image_rgb, image_thermal, bboxes, labels) def augment_val(self, image_rgb, image_thermal, bboxes, labels): if self.use_albumentations and bboxes: try: result = self.val_transform( image=image_rgb, image_thermal=image_thermal, bboxes=bboxes, labels=labels ) return (result['image'], result['image_thermal'], result['bboxes'], result['labels']) except Exception: pass h_orig, w_orig = image_rgb.shape[:2] image_rgb = cv2.resize(image_rgb, (self.input_w, self.input_h)) image_thermal = cv2.resize(image_thermal, (self.input_w, self.input_h)) scale_x = self.input_w / w_orig scale_y = self.input_h / h_orig scaled_bboxes = [[b[0]*scale_x, b[1]*scale_y, b[2]*scale_x, b[3]*scale_y] for b in bboxes] return image_rgb, image_thermal, scaled_bboxes, labels def _manual_augment(self, image_rgb, image_thermal, bboxes, labels): h, w = image_rgb.shape[:2] if np.random.random() < 0.5: image_rgb = np.fliplr(image_rgb).copy() image_thermal = np.fliplr(image_thermal).copy() new_bboxes = [] for box in bboxes: xmin, ymin, xmax, ymax = box new_bboxes.append([w - xmax, ymin, w - xmin, ymax]) bboxes = new_bboxes if np.random.random() < 0.5: shift = np.random.randint(-30, 31) image_thermal = np.clip( image_thermal.astype(np.int16) + shift, 0, 255 ).astype(np.uint8) if np.random.random() < 0.3: noise = np.random.normal(0, 10, image_thermal.shape).astype(np.int16) image_thermal = np.clip( image_thermal.astype(np.int16) + noise, 0, 255 ).astype(np.uint8) image_rgb = cv2.resize(image_rgb, (self.input_w, self.input_h)) image_thermal = cv2.resize(image_thermal, (self.input_w, self.input_h)) scale_x = self.input_w / w scale_y = self.input_h / h scaled_bboxes = [[b[0]*scale_x, b[1]*scale_y, b[2]*scale_x, b[3]*scale_y] for b in bboxes] return image_rgb, image_thermal, scaled_bboxes, labels # ============================================================================ # THERMAL PREPROCESSOR # ============================================================================ class ThermalPreprocessor: """ Complete preprocessing pipeline for dual-modality inputs. Supports both paired and single-modality (CMM) inputs. """ def __init__(self, encoder=None, normalize_method='minmax'): self.encoder = encoder or GridEncoder() self.augmentor = ThermalAugmentor() self.normalize_method = normalize_method def normalize(self, image): """Normalize thermal image.""" if self.normalize_method == 'histogram': return cv2.equalizeHist(image) elif self.normalize_method == 'clahe': clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8)) return clahe.apply(image) else: img_min, img_max = image.min(), image.max() if img_max > img_min: return ((image - img_min) / (img_max - img_min) * 255).astype(np.uint8) return image def process(self, image_rgb, image_thermal, bboxes_pascal, labels, img_size, augment=False, cmm_alpha=0.0): """ Full preprocessing pipeline. Args: image_rgb: (H, W, 3) numpy uint8 RGB image_thermal: (H, W) numpy uint8 single-channel bboxes_pascal: list of [xmin, ymin, xmax, ymax] pixel coords labels: list of class_ids img_size: (h_orig, w_orig) augment: Whether to apply training augmentation cmm_alpha: CMM modality dropout probability (0=disabled) Returns: img_tensor: (4, INPUT_HEIGHT, INPUT_WIDTH) float tensor targets: dict of grid-encoded targets """ h_orig, w_orig = img_size # Normalize thermal image_thermal = self.normalize(image_thermal) # Augment (or just resize) if augment: image_rgb, image_thermal, bboxes_pascal, labels = \ self.augmentor.augment_train( image_rgb, image_thermal, bboxes_pascal, labels, cmm_alpha=cmm_alpha, ) else: image_rgb, image_thermal, bboxes_pascal, labels = \ self.augmentor.augment_val( image_rgb, image_thermal, bboxes_pascal, labels, ) # Ensure correct size if (image_thermal.shape[0] != self.augmentor.input_h or image_thermal.shape[1] != self.augmentor.input_w): image_thermal = cv2.resize( image_thermal, (self.augmentor.input_w, self.augmentor.input_h)) image_rgb = cv2.resize( image_rgb, (self.augmentor.input_w, self.augmentor.input_h)) # Convert to tensor: (4, H, W) tensor_rgb = torch.from_numpy(image_rgb).permute(2, 0, 1).float() / 255.0 tensor_thermal = torch.from_numpy(image_thermal).unsqueeze(0).float() / 255.0 img_tensor = torch.cat([tensor_rgb, tensor_thermal], dim=0) # Convert bboxes to normalized [cx, cy, w, h] boxes_norm = [] valid_labels = [] h_new, w_new = self.augmentor.input_h, self.augmentor.input_w for bbox, lbl in zip(bboxes_pascal, labels): xmin, ymin, xmax, ymax = bbox cx = ((xmin + xmax) / 2) / w_new cy = ((ymin + ymax) / 2) / h_new w = (xmax - xmin) / w_new h = (ymax - ymin) / h_new cx = max(0.0, min(1.0, cx)) cy = max(0.0, min(1.0, cy)) w = max(0.01, min(1.0, w)) h = max(0.01, min(1.0, h)) boxes_norm.append([cx, cy, w, h]) valid_labels.append(lbl) if not boxes_norm: valid_labels = [0] boxes_norm = [] targets = self.encoder.encode(boxes_norm, valid_labels) return img_tensor, targets # ============================================================================ # ANCHOR ANALYSIS UTILITIES # ============================================================================ def analyze_dataset_anchors(dataset, num_anchors=NUM_ANCHORS): """ Run K-Means clustering on dataset bounding boxes to find optimal anchor ratios and sizes. """ from sklearn.cluster import KMeans all_ratios = [] all_sizes = [] print(f"\nAnalyzing dataset (Total Size: {len(dataset)} samples). Sub-sampling for speed...") if hasattr(dataset, 'iter_annotations'): sample_iter = dataset.iter_annotations() elif hasattr(dataset, 'datasets'): # PyTorch ConcatDataset def concat_iter(): import random for ds in dataset.datasets: if hasattr(ds, 'iter_annotations'): yield from ds.iter_annotations() else: indices = random.sample(range(len(ds)), min(500, len(ds))) for i in indices: yield ds[i][1:] sample_iter = concat_iter() else: import random indices = random.sample(range(len(dataset)), min(1000, len(dataset))) sample_iter = (dataset[i][1:] for i in indices) for annotations, (h_orig, w_orig) in sample_iter: for ann in annotations: w = (ann['xmax'] - ann['xmin']) / w_orig h = (ann['ymax'] - ann['ymin']) / h_orig if w > 0.01 and h > 0.01: all_ratios.append(h / w) all_sizes.append(math.sqrt(w * h)) if len(all_ratios) < num_anchors: print(f"Not enough boxes ({len(all_ratios)}). Using defaults.") return DEFAULT_ANCHOR_RATIOS, DEFAULT_ANCHOR_SIZES ratios = np.array(all_ratios) sizes = np.array(all_sizes) print(f" Found {len(ratios)} bounding boxes") print(f" Aspect ratios: min={ratios.min():.2f}, " f"max={ratios.max():.2f}, mean={ratios.mean():.2f}") print(f" Sizes: min={sizes.min():.3f}, " f"max={sizes.max():.3f}, mean={sizes.mean():.3f}") km_ratios = KMeans(n_clusters=num_anchors, random_state=42, n_init=10) km_ratios.fit(ratios.reshape(-1, 1)) optimal_ratios = sorted(km_ratios.cluster_centers_.flatten().tolist()) km_sizes = KMeans(n_clusters=num_anchors, random_state=42, n_init=10) km_sizes.fit(sizes.reshape(-1, 1)) optimal_sizes = sorted(km_sizes.cluster_centers_.flatten().tolist()) print(f"\nOPTIMAL ANCHORS:") print(f" Ratios (h/w): {[f'{r:.3f}' for r in optimal_ratios]}") print(f" Sizes: {[f'{s:.4f}' for s in optimal_sizes]}") return optimal_ratios, optimal_sizes if __name__ == '__main__': print("Preprocessing Module -- Self Test") print("-" * 40) encoder = GridEncoder() preprocessor = ThermalPreprocessor(encoder=encoder) dummy_rgb = np.random.randint(0, 255, (120, 160, 3), dtype=np.uint8) dummy_thermal = np.random.randint(0, 255, (120, 160), dtype=np.uint8) dummy_bboxes = [[30, 20, 80, 100]] dummy_labels = [1] img_tensor, targets = preprocessor.process( image_rgb=dummy_rgb, image_thermal=dummy_thermal, bboxes_pascal=dummy_bboxes, labels=dummy_labels, img_size=(120, 160), augment=False, ) print(f" Input thermal: {dummy_thermal.shape}") print(f" Output tensor: {img_tensor.shape}") print(f" Targets:") for k, v in targets.items(): if isinstance(v, torch.Tensor): print(f" {k}: {v.shape}") else: print(f" {k}: {v}") # Test CMM mode img_cmm, _ = preprocessor.process( image_rgb=dummy_rgb, image_thermal=dummy_thermal, bboxes_pascal=dummy_bboxes, labels=dummy_labels, img_size=(120, 160), augment=True, cmm_alpha=1.0, # Force CMM dropout for test ) print(f" CMM test tensor: {img_cmm.shape}") print("[OK] Preprocessing test passed!")