Create train_mae_swin3d.py

train_mae_swin3d.py

@@ -0,0 +1,755 @@
+#!/usr/bin/env python
+"""
+Masked Autoencoder (MAE) pretraining with 3D Swin Transformer for OPSCC CT scans.
+Asymmetry-aware reconstruction + overfitting monitoring via cosine similarity.
+
+Run example:
+    python train_mae_swin3d.py --data-dir /path/to/your/nii_folder --output-dir ./checkpoints
+"""
+
+"""
+Self-Supervised Learning for OPSCC CT using 3D Swin Transformer MAE
+with asymmetry-aware reconstruction and overfitting monitoring
+"""
+
+import argparse
+import json
+import pickle
+import warnings
+from datetime import datetime
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+
+import numpy as np
+from scipy import ndimage
+import nibabel as nib
+from tqdm import tqdm
+
+warnings.filterwarnings("ignore", category=UserWarning)
+
+
+# ==============================================================================
+# Drop Path
+# ==============================================================================
+
+class DropPath(nn.Module):
+    def __init__(self, drop_prob: float = 0.):
+        super().__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        if self.drop_prob == 0. or not self.training:
+            return x
+        keep_prob = 1 - self.drop_prob
+        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+        random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
+        random_tensor.floor_()
+        return x.div(keep_prob) * random_tensor
+
+
+# ==============================================================================
+# Asymmetry Detectors
+# ==============================================================================
+
+class AirwayAsymmetryDetector:
+    def __init__(self, exclude_inferior_fraction=0.15, exclude_superior_fraction=0.10):
+        self.exclude_inferior_fraction = exclude_inferior_fraction
+        self.exclude_superior_fraction = exclude_superior_fraction
+
+    def find_midline(self, slice_2d):
+        h, w = slice_2d.shape
+        search_range = w // 8
+        center = w // 2
+        best_midline = center
+        best_symmetry = float('inf')
+        for mid in range(center - search_range, center + search_range):
+            compare_width = min(mid, w - mid)
+            if compare_width < 10:
+                continue
+            left = slice_2d[:, mid - compare_width:mid]
+            right = np.flip(slice_2d[:, mid:mid + compare_width], axis=1)
+            diff = np.abs(left - right).mean()
+            if diff < best_symmetry:
+                best_symmetry = diff
+                best_midline = mid
+        return best_midline
+
+    def detect_airway(self, slice_2d, air_thresh=0.1):
+        binary = slice_2d < air_thresh
+        labeled, num_features = ndimage.label(binary)
+        edge_labels = set(labeled[0,:].flatten()) | set(labeled[-1,:].flatten()) | \
+                      set(labeled[:,0].flatten()) | set(labeled[:,-1].flatten())
+        airway_mask = np.zeros_like(binary)
+        for label_id in range(1, num_features + 1):
+            if label_id not in edge_labels:
+                component = labeled == label_id
+                if component.sum() > 20:
+                    airway_mask |= component
+        return airway_mask
+
+    def forward(self, volume):
+        d, h, w = volume.shape
+        inferior_cutoff = int(d * self.exclude_inferior_fraction)
+        superior_cutoff = int(d * (1 - self.exclude_superior_fraction))
+        results = {'effacement': [], 'mass_effect': [], 'midline_shift': [], 'hybrid': [], 'midlines': []}
+        for z in range(d):
+            slice_2d = volume[z]
+            midline = self.find_midline(slice_2d)
+            midline_shift = midline - w // 2
+            results['midlines'].append(midline)
+            airway_mask = self.detect_airway(slice_2d)
+            left_air = airway_mask[:, :midline].sum()
+            right_air = airway_mask[:, midline:].sum()
+            total = left_air + right_air
+            effacement = abs(left_air - right_air) / max(total, 1) if total > 0 else 0
+            compare_width = min(midline, w - midline)
+            mass_effect = 0
+            if compare_width > 0:
+                soft_tissue = (slice_2d > 0.2) & (slice_2d < 0.7)
+                left = slice_2d[:, midline-compare_width:midline] * soft_tissue[:, midline-compare_width:midline]
+                right = np.flip(slice_2d[:, midline:midline+compare_width], axis=1) * np.flip(soft_tissue[:, midline:midline+compare_width], axis=1)
+                mass_effect = np.abs(left - right).mean()
+            in_range = inferior_cutoff <= z <= superior_cutoff
+            hybrid = (0.5 * effacement + 0.5 * mass_effect) if in_range else 0
+            results['effacement'].append(effacement)
+            results['mass_effect'].append(mass_effect)
+            results['midline_shift'].append(midline_shift)
+            results['hybrid'].append(hybrid)
+        return {k: np.array(v) for k, v in results.items()}
+
+
+class GlobalSoftTissueAsymmetryDetector:
+    def __init__(self, exclude_inferior_fraction=0.15, exclude_superior_fraction=0.10):
+        self.exclude_inferior_fraction = exclude_inferior_fraction
+        self.exclude_superior_fraction = exclude_superior_fraction
+
+    def forward(self, volume, midlines=None):
+        d, h, w = volume.shape
+        if midlines is None:
+            midlines = [w // 2] * d
+        results = {'left_hypo': [], 'right_hypo': [], 'hypo_asymmetry': []}
+        for z in range(d):
+            slice_2d = volume[z]
+            midline = midlines[z]
+            soft_tissue = (slice_2d > 0.2) & (slice_2d < 0.7)
+            hypodense = (slice_2d < 0.35) & soft_tissue
+            hypodense = ndimage.binary_opening(hypodense, iterations=1)
+            hypodense = ndimage.binary_closing(hypodense, iterations=2)
+            labeled, num_features = ndimage.label(hypodense)
+            left_count = right_count = 0
+            for i in range(1, num_features + 1):
+                region = labeled == i
+                size = region.sum()
+                if 10 < size < 150:
+                    centroid_x = np.argwhere(region)[:,1].mean()
+                    if centroid_x < midline:
+                        left_count += 1
+                    else:
+                        right_count += 1
+            results['left_hypo'].append(left_count)
+            results['right_hypo'].append(right_count)
+            results['hypo_asymmetry'].append(abs(left_count - right_count))
+        return {k: np.array(v) for k, v in results.items()}
+
+
+# ==============================================================================
+# 3D Swin Transformer Components
+# ==============================================================================
+
+def window_partition3d(x, window_size=(4,4,4)):
+    B, C, D, H, W = x.shape
+    ws_d, ws_h, ws_w = window_size
+    pad_d = (ws_d - D % ws_d) % ws_d
+    pad_h = (ws_h - H % ws_h) % ws_h
+    pad_w = (ws_w - W % ws_w) % ws_w
+    x = F.pad(x, (0, pad_w, 0, pad_h, 0, pad_d))
+    Dp, Hp, Wp = D + pad_d, H + pad_h, W + pad_w
+    x = x.reshape(B, C, Dp // ws_d, ws_d, Hp // ws_h, ws_h, Wp // ws_w, ws_w)
+    x = x.permute(0, 2, 4, 6, 1, 3, 5, 7).contiguous()
+    windows = x.reshape(-1, C, ws_d * ws_h * ws_w).permute(0, 2, 1).contiguous()
+    return windows, (pad_d, pad_h, pad_w)
+
+
+def window_reverse3d(windows, window_size, B, D, H, W, pads):
+    pad_d, pad_h, pad_w = pads
+    ws_d, ws_h, ws_w = window_size
+    Dp, Hp, Wp = D + pad_d, H + pad_h, W + pad_w
+    x = windows.reshape(B, Dp // ws_d, Hp // ws_h, Wp // ws_w, ws_d, ws_h, ws_w, -1)
+    x = x.permute(0, 7, 1, 4, 2, 5, 3, 6).contiguous()
+    x = x.reshape(B, -1, Dp, Hp, Wp)
+    x = x[:, :, :D, :H, :W]
+    return x
+
+
+class WindowAttention3D(nn.Module):
+    def __init__(self, dim, window_size=(4,4,4), num_heads=3, qkv_bias=True, qk_scale=None,
+                 attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.dim = dim
+        self.window_size = window_size
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        coords_d = torch.arange(window_size[0])
+        coords_h = torch.arange(window_size[1])
+        coords_w = torch.arange(window_size[2])
+        coords = torch.stack(torch.meshgrid(coords_d, coords_h, coords_w, indexing='ij'))
+        coords_flatten = torch.flatten(coords, 1)
+        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]
+        relative_coords = relative_coords.permute(1, 2, 0).contiguous()
+
+        relative_coords[:, :, 0] += window_size[0] - 1
+        relative_coords[:, :, 1] += window_size[1] - 1
+        relative_coords[:, :, 2] += window_size[2] - 1
+
+        relative_coords[:, :, 0] *= (2 * window_size[1] - 1) * (2 * window_size[2] - 1)
+        relative_coords[:, :, 1] *= (2 * window_size[2] - 1)
+        self.relative_position_index = relative_coords.sum(-1)
+
+        max_rel_pos = self.relative_position_index.max().item()
+        self.relative_position_bias_table = nn.Parameter(torch.zeros((max_rel_pos + 1, num_heads)))
+        nn.init.trunc_normal_(self.relative_position_bias_table, std=.02)
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x, mask=None):
+        B_, N, C = x.shape
+        rel_index = self.relative_position_index[:N, :N]
+        relative_position_bias = self.relative_position_bias_table[rel_index.view(-1)]
+        relative_position_bias = relative_position_bias.view(N, N, -1).permute(2, 0, 1).contiguous()
+
+        qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+        q = q * self.scale
+        attn = (q @ k.transpose(-2, -1))
+        attn = attn + relative_position_bias.unsqueeze(0)
+
+        if mask is not None:
+            nW = mask.shape[0]
+            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
+            attn = attn.view(-1, self.num_heads, N, N)
+
+        attn = self.softmax(attn)
+        attn = self.attn_drop(attn)
+        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class SwinTransformerBlock3D(nn.Module):
+    def __init__(self, dim, num_heads, window_size=(4,4,4), shift_size=(0,0,0),
+                 mlp_ratio=4., qkv_bias=True, drop=0., attn_drop=0., drop_path=0.,
+                 act_layer=nn.GELU, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.dim = dim
+        self.window_size = window_size
+        self.shift_size = shift_size
+        self.norm1 = norm_layer(dim)
+        self.attn = WindowAttention3D(dim=dim, window_size=window_size, num_heads=num_heads,
+                                       qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop)
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = nn.Sequential(
+            nn.Linear(dim, mlp_hidden_dim), act_layer(), nn.Dropout(drop),
+            nn.Linear(mlp_hidden_dim, dim), nn.Dropout(drop)
+        )
+
+    def forward(self, x):
+        shortcut = x
+        x_norm = x.permute(0, 2, 3, 4, 1)
+        x_norm = self.norm1(x_norm)
+        x = x_norm.permute(0, 4, 1, 2, 3)
+
+        windows, pads = window_partition3d(x, self.window_size)
+        attn_windows = self.attn(windows)
+        x = window_reverse3d(attn_windows, self.window_size, x.shape[0], x.shape[2], x.shape[3], x.shape[4], pads)
+
+        x = shortcut + self.drop_path(x)
+
+        x_norm = x.permute(0, 2, 3, 4, 1)
+        x_norm = self.norm2(x_norm)
+        x_norm = x_norm.permute(0, 4, 1, 2, 3)
+        x_mlp = self.mlp(x_norm.permute(0, 2, 3, 4, 1)).permute(0, 4, 1, 2, 3)
+        x = x + self.drop_path(x_mlp)
+        return x
+
+
+class PatchEmbed3D(nn.Module):
+    def __init__(self, patch_size=(4,4,4), in_chans=1, embed_dim=96):
+        super().__init__()
+        self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+
+    def forward(self, x):
+        return self.proj(x)
+
+
+class PatchMerging3D(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.reduction = nn.Linear(8 * dim, 2 * dim, bias=False)
+
+    def forward(self, x):
+        B, C, D, H, W = x.shape
+        pad_d, pad_h, pad_w = D % 2, H % 2, W % 2
+        if pad_d or pad_h or pad_w:
+            x = F.pad(x, (0, pad_w, 0, pad_h, 0, pad_d))
+        _, _, Dp, Hp, Wp = x.shape
+        x = x.permute(0, 2, 3, 4, 1)
+        x = x.view(B, Dp // 2, 2, Hp // 2, 2, Wp // 2, 2, C)
+        x = x.permute(0, 1, 3, 5, 2, 4, 6, 7).contiguous()
+        x = x.view(B, Dp // 2, Hp // 2, Wp // 2, 8 * C)
+        x = self.reduction(x)
+        x = x.permute(0, 4, 1, 2, 3).contiguous()
+        return x
+
+
+class SwinTransformer3D(nn.Module):
+    def __init__(self, in_chans=1, embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24],
+                 window_size=(4,4,4), mlp_ratio=4., qkv_bias=True, drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1):
+        super().__init__()
+        self.patch_embed = PatchEmbed3D(in_chans=in_chans, embed_dim=embed_dim)
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
+        self.layers = nn.ModuleList()
+        dim = embed_dim
+        for i_layer in range(len(depths)):
+            blocks = nn.ModuleList([
+                SwinTransformerBlock3D(dim=dim, num_heads=num_heads[i_layer], window_size=window_size,
+                                       drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i])
+                for i in range(depths[i_layer])
+            ])
+            self.layers.append(blocks)
+            if i_layer < len(depths)-1:
+                self.layers.append(PatchMerging3D(dim))
+                dim *= 2
+        self.norm = nn.LayerNorm(dim)
+        self.avgpool = nn.AdaptiveAvgPool3d(1)
+        self.feature_dim = dim
+
+    def forward(self, x):
+        x = self.patch_embed(x)
+        for layer in self.layers:
+            if isinstance(layer, PatchMerging3D):
+                x = layer(x)
+            else:
+                for blk in layer:
+                    x = blk(x)
+        x = self.avgpool(x).flatten(1)
+        x = self.norm(x)
+        return x
+
+
+# ==============================================================================
+# MAE Model
+# ==============================================================================
+
+class MAE_Swin3D(nn.Module):
+    def __init__(self, input_shape=(60, 128, 128)):
+        super().__init__()
+        self.input_shape = input_shape
+        self.encoder = SwinTransformer3D(in_chans=1)
+        decoder_dim = 512
+        self.decoder = nn.Sequential(
+            nn.Linear(self.encoder.feature_dim, decoder_dim),
+            nn.ReLU(),
+            nn.Linear(decoder_dim, np.prod(input_shape))
+        )
+        self.airway_head = nn.Linear(self.encoder.feature_dim, 4 * input_shape[0])
+        self.lymph_head = nn.Linear(self.encoder.feature_dim, 3 * input_shape[0])
+
+    def forward(self, x):
+        feat = self.encoder(x)
+        recon_flat = self.decoder(feat)
+        recon = recon_flat.view(-1, 1, *self.input_shape)
+        airway_pred = self.airway_head(feat).view(-1, self.input_shape[0], 4)
+        lymph_pred = self.lymph_head(feat).view(-1, self.input_shape[0], 3)
+        return {
+            'reconstruction': recon,
+            'airway_pred': airway_pred,
+            'lymph_pred': lymph_pred,
+            'features': feat
+        }
+
+
+# ==============================================================================
+# Augmentations
+# ==============================================================================
+
+def augment_volume(volume):
+    aug = volume.clone()
+    device = aug.device
+
+    if torch.rand(1) > 0.3:
+        shift = (torch.rand(1).to(device) - 0.5) * 0.4
+        aug += shift
+
+    if torch.rand(1) > 0.3:
+        scale = 0.7 + torch.rand(1).to(device) * 0.6
+        aug *= scale
+
+    if torch.rand(1) > 0.3:
+        noise = torch.randn_like(aug) * 0.1
+        aug += noise
+
+    if torch.rand(1) > 0.5:
+        aug = torch.flip(aug, dims=[-1])
+
+    if torch.rand(1) > 0.5:
+        aug = torch.flip(aug, dims=[-2])
+
+    if torch.rand(1) > 0.7:
+        k = torch.randint(1, 4, (1,)).item()
+        aug = torch.rot90(aug, k, dims=[-2, -1])
+
+    if torch.rand(1) > 0.5:
+        _, _, D, H, W = aug.shape
+        crop_d = int(D * (0.80 + torch.rand(1).item() * 0.15))
+        crop_h = int(H * (0.80 + torch.rand(1).item() * 0.15))
+        crop_w = int(W * (0.80 + torch.rand(1).item() * 0.15))
+        start_d = torch.randint(0, D - crop_d + 1, (1,)).item()
+        start_h = torch.randint(0, H - crop_h + 1, (1,)).item()
+        start_w = torch.randint(0, W - crop_w + 1, (1,)).item()
+        aug = aug[:, :, start_d:start_d+crop_d, start_h:start_h+crop_h, start_w:start_w+crop_w]
+        aug = F.interpolate(aug, size=(D, H, W), mode='trilinear', align_corners=False)
+
+    if torch.rand(1) > 0.7:
+        kernel_size = 3
+        padding = kernel_size // 2
+        aug = F.avg_pool3d(aug, kernel_size=kernel_size, stride=1, padding=padding)
+
+    if torch.rand(1) > 0.7:
+        _, _, D, H, W = aug.shape
+        erase_d = int(D * (0.05 + torch.rand(1).item() * 0.10))
+        erase_h = int(H * (0.05 + torch.rand(1).item() * 0.10))
+        erase_w = int(W * (0.05 + torch.rand(1).item() * 0.10))
+        start_d = torch.randint(0, D - erase_d + 1, (1,)).item()
+        start_h = torch.randint(0, H - erase_h + 1, (1,)).item()
+        start_w = torch.randint(0, W - erase_w + 1, (1,)).item()
+        aug[:, :, start_d:start_d+erase_d, start_h:start_h+erase_h, start_w:start_w+erase_w] = aug.mean()
+
+    aug = torch.clamp(aug, 0, 1)
+    return aug
+
+
+# ==============================================================================
+# Dataset
+# ==============================================================================
+
+class OPSCCDataset(Dataset):
+    def __init__(self, data_dir: str, cache_asymmetry: bool = True):
+        self.data_dir = Path(data_dir)
+        self.volume_paths = list(self.data_dir.glob("**/cropped_volume.nii.gz"))
+        print(f"Found {len(self.volume_paths)} volumes")
+
+        self.cache_file = self.data_dir / ".asymmetry_cache.pkl"
+        self.cache_asymmetry = cache_asymmetry
+        self.asymmetry_cache = {}
+        self.airway_detector = AirwayAsymmetryDetector()
+        self.lymphnode_detector = GlobalSoftTissueAsymmetryDetector()
+
+        if self.cache_asymmetry:
+            if self.cache_file.is_file():
+                try:
+                    with open(self.cache_file, 'rb') as f:
+                        self.asymmetry_cache = pickle.load(f)
+                    print(f"Loaded asymmetry cache ({len(self.asymmetry_cache)} entries)")
+                except Exception:
+                    print("Cache load failed → recomputing")
+                    self._precompute_asymmetry()
+            else:
+                print("Computing asymmetry metrics...")
+                self._precompute_asymmetry()
+                try:
+                    with open(self.cache_file, 'wb') as f:
+                        pickle.dump(self.asymmetry_cache, f)
+                    print("Cache saved")
+                except Exception as e:
+                    print(f"Cache save failed: {e}")
+
+    def _precompute_asymmetry(self):
+        for idx, path in enumerate(tqdm(self.volume_paths, desc="Asymmetry")):
+            volume = self._load_volume(path)
+            metrics = self._compute_asymmetry(volume)
+            self.asymmetry_cache[idx] = metrics
+
+    def _load_volume(self, path: Path) -> np.ndarray:
+        img = nib.load(str(path))
+        volume = img.get_fdata().astype(np.float32)
+        if volume.ndim == 3 and volume.shape[2] < volume.shape[0]:
+            volume = np.transpose(volume, (2, 0, 1))
+        return volume
+
+    def _compute_asymmetry(self, volume: np.ndarray) -> dict:
+        airway = self.airway_detector.forward(volume)
+        lymphnode = self.lymphnode_detector.forward(volume, airway['midlines'].tolist())
+        return {'airway': airway, 'lymphnode': lymphnode}
+
+    def __len__(self) -> int:
+        return len(self.volume_paths)
+
+    def __getitem__(self, idx: int) -> dict:
+        path = self.volume_paths[idx]
+        volume = self._load_volume(path)
+
+        if self.cache_asymmetry and idx in self.asymmetry_cache:
+            metrics = self.asymmetry_cache[idx]
+        else:
+            metrics = self._compute_asymmetry(volume)
+
+        airway_tensor = np.stack([
+            metrics['airway']['effacement'],
+            metrics['airway']['mass_effect'],
+            metrics['airway']['midline_shift'],
+            metrics['airway']['hybrid']
+        ], axis=0)
+
+        lymph_tensor = np.stack([
+            metrics['lymphnode']['left_hypo'],
+            metrics['lymphnode']['right_hypo'],
+            metrics['lymphnode']['hypo_asymmetry']
+        ], axis=0)
+
+        return {
+            'volume': torch.from_numpy(volume).unsqueeze(0).float(),
+            'airway_metrics': torch.from_numpy(airway_tensor).float(),
+            'lymphnode_metrics': torch.from_numpy(lymph_tensor).float(),
+        }
+
+
+# ==============================================================================
+# Loss
+# ==============================================================================
+
+class MAEAsymmetryLoss(nn.Module):
+    def __init__(self, mask_ratio=0.75, asymmetry_boost=5.0):
+        super().__init__()
+        self.mse = nn.MSELoss(reduction='none')
+        self.mask_ratio = mask_ratio
+        self.asymmetry_boost = asymmetry_boost
+
+    def forward(self, outputs, batch):
+        recon = outputs['reconstruction']
+        target = batch['volume']
+
+        B, C, D, H, W = target.shape
+        num_patches = D * H * W
+        mask = torch.rand(B, num_patches, device=target.device) < self.mask_ratio
+        mask = mask.view(B, 1, D, H, W).expand_as(recon)
+
+        diff = self.mse(recon, target) * mask.float()
+
+        hybrid = batch['airway_metrics'][:, 3, :]
+        hybrid_norm = hybrid / (hybrid.max(dim=1, keepdim=True)[0] + 1e-6)
+        slice_weights = 1.0 + self.asymmetry_boost * hybrid_norm
+        weights = slice_weights.unsqueeze(1).unsqueeze(3).unsqueeze(4).expand_as(diff)
+
+        recon_loss = (diff * weights).sum() / (mask.sum() + 1e-6)
+
+        airway_loss = F.mse_loss(outputs['airway_pred'], batch['airway_metrics'].permute(0, 2, 1))
+        lymph_loss = F.mse_loss(outputs['lymph_pred'], batch['lymphnode_metrics'].permute(0, 2, 1))
+
+        return recon_loss + airway_loss + lymph_loss
+
+
+# ==============================================================================
+# Trainer
+# ==============================================================================
+
+class TrainerWithMonitoring:
+    def __init__(self, model, train_loader, device, lr=1e-4, output_dir=None):
+        self.model = model.to(device)
+        self.device = device
+        self.train_loader = train_loader
+        self.optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
+        self.loss_fn = MAEAsymmetryLoss()
+
+        self.output_dir = Path(output_dir) if output_dir else None
+        if self.output_dir:
+            self.output_dir.mkdir(parents=True, exist_ok=True)
+
+        self.history = {
+            'epoch': [],
+            'loss': [],
+            'cosine_sim_mean': [],
+            'cosine_sim_std': [],
+        }
+
+    def compute_cosine_similarity(self, n_samples=50):
+        self.model.eval()
+        similarities = []
+        with torch.no_grad():
+            for i, batch in enumerate(self.train_loader):
+                if i >= n_samples:
+                    break
+                volume = batch['volume'].to(self.device)
+                feat1 = self.model.encoder(volume)
+                volume_aug = augment_volume(volume)
+                feat2 = self.model.encoder(volume_aug)
+                feat1_norm = F.normalize(feat1, dim=1)
+                feat2_norm = F.normalize(feat2, dim=1)
+                sim = (feat1_norm * feat2_norm).sum(dim=1)
+                similarities.extend(sim.cpu().numpy().tolist())
+        self.model.train()
+        return np.mean(similarities), np.std(similarities)
+
+    def save_checkpoint(self, epoch, is_best=False):
+        if not self.output_dir:
+            return
+        path = self.output_dir / f"checkpoint_epoch_{epoch:03d}.pt"
+        torch.save({
+            'epoch': epoch,
+            'model_state_dict': self.model.state_dict(),
+            'optimizer_state_dict': self.optimizer.state_dict(),
+            'history': self.history,
+        }, path)
+        print(f"Checkpoint saved: {path.name}")
+
+        if is_best:
+            best_path = self.output_dir / "best_model.pt"
+            torch.save(self.model.state_dict(), best_path)
+            print(f"Best model updated: {best_path.name}")
+
+    def train(self, n_epochs=100, monitor_every=5, save_every=10,
+              early_stop_patience=20, early_stop_after=30):
+        best_loss = float('inf')
+        patience_counter = 0
+        best_epoch = 0
+
+        for epoch in range(1, n_epochs + 1):
+            self.model.train()
+            total_loss = 0.0
+            num_batches = 0
+
+            for batch in tqdm(self.train_loader, desc=f"Epoch {epoch}", leave=False):
+                volume = batch['volume'].to(self.device)
+                airway_metrics = batch['airway_metrics'].to(self.device)
+                lymphnode_metrics = batch['lymphnode_metrics'].to(self.device)
+
+                self.optimizer.zero_grad()
+                outputs = self.model(volume)
+
+                loss = self.loss_fn(outputs, batch)
+
+                loss.backward()
+                self.optimizer.step()
+
+                total_loss += loss.item()
+                num_batches += 1
+
+            avg_loss = total_loss / num_batches if num_batches > 0 else 0.0
+
+            is_best = avg_loss < best_loss
+            if is_best:
+                best_loss = avg_loss
+                best_epoch = epoch
+                patience_counter = 0
+            else:
+                patience_counter += 1
+
+            if epoch % monitor_every == 0 or epoch == 1:
+                cos_mean, cos_std = self.compute_cosine_similarity()
+                self.history['epoch'].append(epoch)
+                self.history['loss'].append(avg_loss)
+                self.history['cosine_sim_mean'].append(cos_mean)
+                self.history['cosine_sim_std'].append(cos_std)
+
+                msg = f"Epoch {epoch:3d} | Loss: {avg_loss:.4f} | CosSim: {cos_mean:.3f}±{cos_std:.3f}"
+                if is_best:
+                    msg += " ★"
+                print(msg)
+
+                if cos_mean > 0.95:
+                    print(f"  WARNING: Cosine similarity very high ({cos_mean:.3f}) — possible collapse")
+
+            else:
+                msg = f"Epoch {epoch:3d} | Loss: {avg_loss:.4f}"
+                if is_best:
+                    msg += " ★"
+                print(msg)
+
+            if epoch % save_every == 0:
+                self.save_checkpoint(epoch, is_best=is_best)
+            elif is_best:
+                self.save_checkpoint(epoch, is_best=True)
+
+            if epoch > early_stop_after and patience_counter >= early_stop_patience:
+                print(f"Early stopping at epoch {epoch}")
+                break
+
+        if self.output_dir:
+            torch.save(self.model.state_dict(), self.output_dir / "final_model.pt")
+            with open(self.output_dir / "history.json", 'w') as f:
+                json.dump(self.history, f, indent=2)
+
+        print(f"Best loss: {best_loss:.4f} at epoch {best_epoch}")
+        return self.history
+
+
+# ==============================================================================
+# Main
+# ==============================================================================
+
+def main():
+    parser = argparse.ArgumentParser(description="3D Swin MAE pretraining")
+    parser.add_argument("--data-dir",      type=str, required=True,   help="Folder containing cropped_volume.nii.gz files")
+    parser.add_argument("--output-dir",    type=str, default="./checkpoints", help="Folder to save models and logs")
+    parser.add_argument("--batch-size",    type=int, default=2)
+    parser.add_argument("--epochs",        type=int, default=100)
+    parser.add_argument("--lr",            type=float, default=1e-4)
+    parser.add_argument("--monitor-every", type=int, default=5)
+    parser.add_argument("--save-every",    type=int, default=10)
+    parser.add_argument("--patience",      type=int, default=20)
+    parser.add_argument("--early-after",   type=int, default=30)
+    parser.add_argument("--no-cache",      action="store_true")
+
+    args = parser.parse_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Device: {device}")
+
+    dataset = OPSCCDataset(
+        data_dir=args.data_dir,
+        cache_asymmetry=not args.no_cache
+    )
+
+    loader = DataLoader(
+        dataset,
+        batch_size=args.batch_size,
+        shuffle=True,
+        num_workers=0,
+        pin_memory=device.type == "cuda"
+    )
+
+    model = MAE_Swin3D()
+
+    trainer = TrainerWithMonitoring(
+        model=model,
+        train_loader=loader,
+        device=device,
+        lr=args.lr,
+        output_dir=args.output_dir
+    )
+
+    trainer.train(
+        n_epochs=args.epochs,
+        monitor_every=args.monitor_every,
+        save_every=args.save_every,
+        early_stop_patience=args.patience,
+        early_stop_after=args.early_after
+    )
+
+    print("\nNote: Volumes are expected to be cropped, resized to ~60×128×128, intensities [0,1].")
+
+
+if __name__ == "__main__":
+    main()