| """
|
| test_robustness.py
|
| Robustness evaluation for student models (Phase 2, Steps 2 & 3).
|
|
|
| Tests trained models under various input corruptions and reports
|
| localization error in meters (median, P90, P99).
|
|
|
| Step 2 — AP Dropout:
|
| - Remove each AP individually (zero out 1 of 4 channels)
|
| - Random 1-AP dropout (random channel zeroed per sample)
|
| - Random 2-AP dropout (2 random channels zeroed per sample)
|
|
|
| Step 3 — Data Augmentation Robustness (optional):
|
| - Gaussian noise (sigma=0.1, 0.2)
|
| - AP attenuation (scale channel by 0.5)
|
| - Phase perturbation (random rotation in complex-like space)
|
| - Heatmap blur (Gaussian blur on input)
|
|
|
| Usage:
|
| python test_robustness.py
|
| Edit CONFIG below to select model and weights path.
|
| """
|
| import os
|
| import sys
|
| import torch
|
| import torch.nn as nn
|
| import torch.nn.functional as torchF
|
| import numpy as np
|
|
|
| from utils import localization_error
|
| from data_loader import load_data
|
|
|
|
|
|
|
|
|
| MODEL_TYPE = 'mobilenet'
|
| WEIGHTS_PATH = './weights/mobilenet_standalone_best.pth'
|
| BATCH_SIZE = 32
|
|
|
|
|
| testpath = [
|
| './data/dataset_fov_test_jacobs_July28_2.mat',
|
| './data/dataset_non_fov_test_jacobs_July28_2.mat',
|
| ]
|
|
|
|
|
|
|
|
|
| def no_corruption(inputs):
|
| """Baseline: no corruption."""
|
| return inputs
|
|
|
| def drop_ap(inputs, ap_index):
|
| """Zero out a specific AP channel."""
|
| corrupted = inputs.clone()
|
| corrupted[:, ap_index, :, :] = 0.0
|
| return corrupted
|
|
|
| def random_drop_1ap(inputs):
|
| """Randomly zero 1 AP per sample."""
|
| corrupted = inputs.clone()
|
| for i in range(corrupted.shape[0]):
|
| ap = torch.randint(0, 4, (1,)).item()
|
| corrupted[i, ap, :, :] = 0.0
|
| return corrupted
|
|
|
| def random_drop_2ap(inputs):
|
| """Randomly zero 2 APs per sample."""
|
| corrupted = inputs.clone()
|
| for i in range(corrupted.shape[0]):
|
| aps = torch.randperm(4)[:2]
|
| for ap in aps:
|
| corrupted[i, ap, :, :] = 0.0
|
| return corrupted
|
|
|
| def gaussian_noise(inputs, sigma=0.1):
|
| """Add Gaussian noise."""
|
| return inputs + torch.randn_like(inputs) * sigma
|
|
|
| def ap_attenuation(inputs, scale=0.5):
|
| """Scale all channels by a factor (simulates weaker signal)."""
|
| return inputs * scale
|
|
|
| def single_ap_attenuation(inputs, ap_index, scale=0.3):
|
| """Attenuate a single AP channel."""
|
| corrupted = inputs.clone()
|
| corrupted[:, ap_index, :, :] *= scale
|
| return corrupted
|
|
|
| def gaussian_blur(inputs, kernel_size=5, sigma=1.0):
|
| """Apply Gaussian blur to each channel."""
|
|
|
| x = torch.arange(kernel_size, dtype=torch.float32, device=inputs.device) - kernel_size // 2
|
| kernel_1d = torch.exp(-0.5 * (x / sigma) ** 2)
|
| kernel_1d = kernel_1d / kernel_1d.sum()
|
|
|
| kernel_2d = kernel_1d[:, None] * kernel_1d[None, :]
|
| kernel_2d = kernel_2d.expand(inputs.shape[1], 1, kernel_size, kernel_size)
|
| pad = kernel_size // 2
|
| return torchF.conv2d(inputs, kernel_2d, padding=pad, groups=inputs.shape[1])
|
|
|
|
|
|
|
|
|
| def load_model(model_type, weights_path, device):
|
| if model_type == 'mobilenet':
|
| from student_mobilenet import MobileNetStudent
|
| model = MobileNetStudent(input_nc=4)
|
| elif model_type == 'tinycnn':
|
| from student_tinycnn import TinyCNNStudent
|
| model = TinyCNNStudent(input_nc=4)
|
| elif model_type == 'unet':
|
| from student_unet import MobileNetUNetStudent
|
| model = MobileNetUNetStudent(input_nc=4)
|
| else:
|
| raise ValueError(f'Unknown model type: {model_type}')
|
|
|
| model.load_state_dict(torch.load(weights_path, map_location=device, weights_only=True))
|
| model.to(device)
|
| model.eval()
|
| return model
|
|
|
|
|
| def evaluate_condition(model, test_loader, corrupt_fn, device):
|
| """Run model on corrupted inputs, return error stats in meters."""
|
| all_errors = []
|
| with torch.no_grad():
|
| for data in test_loader:
|
| inputs = data[0].to(device)
|
| labels = data[1].to(device)
|
|
|
| corrupted = corrupt_fn(inputs)
|
| output = model(corrupted)
|
|
|
| errors = localization_error(
|
| output.cpu().numpy(),
|
| labels.cpu().numpy(),
|
| scale=0.1,
|
| )
|
| all_errors.extend(errors)
|
|
|
| all_errors = np.array(all_errors)
|
| return {
|
| 'median': np.median(all_errors),
|
| 'p90': np.percentile(all_errors, 90),
|
| 'p99': np.percentile(all_errors, 99),
|
| 'mean': np.mean(all_errors),
|
| }
|
|
|
|
|
| def main():
|
| device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
|
| print(f'Device: {device}')
|
| print(f'Model: {MODEL_TYPE}, Weights: {WEIGHTS_PATH}')
|
|
|
|
|
| model = load_model(MODEL_TYPE, WEIGHTS_PATH, device)
|
| n_params = sum(p.numel() for p in model.parameters())
|
| print(f'Parameters: {n_params:,}')
|
|
|
|
|
| print('Loading test data...')
|
| test_w_list, test_lbl_list = [], []
|
| for fp in testpath:
|
| _, w, lbl = load_data(fp)
|
| if lbl.dim() == 3:
|
| lbl = lbl.unsqueeze(1)
|
| test_w_list.append(w)
|
| test_lbl_list.append(lbl)
|
| test_w = torch.cat(test_w_list, dim=0)
|
| test_lbl = torch.cat(test_lbl_list, dim=0)
|
| print(f'Test samples: {test_w.shape[0]}')
|
|
|
| test_data = torch.utils.data.TensorDataset(test_w, test_lbl)
|
| test_loader = torch.utils.data.DataLoader(
|
| test_data, batch_size=BATCH_SIZE, shuffle=False,
|
| num_workers=4, pin_memory=True,
|
| )
|
|
|
|
|
| conditions = {
|
|
|
| 'Normal (no corruption)': no_corruption,
|
| 'Drop AP 0': lambda x: drop_ap(x, 0),
|
| 'Drop AP 1': lambda x: drop_ap(x, 1),
|
| 'Drop AP 2': lambda x: drop_ap(x, 2),
|
| 'Drop AP 3': lambda x: drop_ap(x, 3),
|
| 'Random drop 1 AP': random_drop_1ap,
|
| 'Random drop 2 APs': random_drop_2ap,
|
|
|
| 'Gaussian noise (sigma=0.1)': lambda x: gaussian_noise(x, sigma=0.1),
|
| 'Gaussian noise (sigma=0.2)': lambda x: gaussian_noise(x, sigma=0.2),
|
| 'Attenuation (0.5x all)': lambda x: ap_attenuation(x, scale=0.5),
|
| 'Attenuate AP 0 (0.3x)': lambda x: single_ap_attenuation(x, 0, 0.3),
|
| 'Gaussian blur (k=5, s=1)': lambda x: gaussian_blur(x, kernel_size=5, sigma=1.0),
|
| }
|
|
|
|
|
| print(f'\n{"="*75}')
|
| print(f' ROBUSTNESS EVALUATION — {MODEL_TYPE.upper()}')
|
| print(f'{"="*75}')
|
| print(f'{"Condition":<30} {"Median":>8} {"P90":>8} {"P99":>8} {"Mean":>8}')
|
| print(f'{"-"*75}')
|
|
|
| results = {}
|
| for name, corrupt_fn in conditions.items():
|
| stats = evaluate_condition(model, test_loader, corrupt_fn, device)
|
| results[name] = stats
|
| print(f'{name:<30} {stats["median"]:>7.3f}m {stats["p90"]:>7.3f}m {stats["p99"]:>7.3f}m {stats["mean"]:>7.3f}m')
|
|
|
| print(f'{"="*75}')
|
|
|
|
|
| save_path = f'robustness_{MODEL_TYPE}.txt'
|
| with open(save_path, 'w') as f:
|
| f.write(f'Model: {MODEL_TYPE}\n')
|
| f.write(f'Weights: {WEIGHTS_PATH}\n')
|
| f.write(f'Test samples: {test_w.shape[0]}\n\n')
|
| f.write(f'{"Condition":<30} {"Median":>8} {"P90":>8} {"P99":>8} {"Mean":>8}\n')
|
| f.write(f'{"-"*62}\n')
|
| for name, stats in results.items():
|
| f.write(f'{name:<30} {stats["median"]:>7.3f}m {stats["p90"]:>7.3f}m {stats["p99"]:>7.3f}m {stats["mean"]:>7.3f}m\n')
|
| print(f'\nResults saved to {save_path}')
|
|
|
|
|
| if __name__ == '__main__':
|
| main()
|
|
|
