"""Shared helpers for the MONAI iterative pseudo-labeling pipeline (scripts_monai/). Backbone: MONAI 1.6.x (pip) — monai.networks.nets.UNet, MONAI transforms / Dataset / DataLoader, and MONAI sliding-window inference. Everything study-specific (splits, gates, loss weighting, longitudinal consistency) is shared with the other pipelines so round results stay directly comparable: * per-sample weighted Dice+CE loss — Pipeline A's convention (row "w"; pseudo-label rows down-weighted sample-by-sample) * selection gates + head-relative consistency filter — Pipeline B's implementation (scripts_nnu/nnu_common; absolute patient-space resampling is unreliable across acquisitions, see README) * DDP via torchrun for training, per-rank row sharding for inference — Pipeline A's parallelism model Network: monai UNet 3D, 1 in / 2 out, channels 16→128 (~1.2M params) — the MONAI analogue of Pipeline A's Unet3D(base=16, depth=4). """ import os import sys # Data/run dirs: overridable for scratch runs. Code (src/, scripts_nnu/): next to this file. ROOT = os.environ.get( "LONGITUDINAL_ROOT", os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) _REPO = ROOT sys.path.insert(0, _REPO) sys.path.insert(0, os.path.join(_REPO, "scripts_nnu")) import numpy as np import torch import torch.nn as nn from monai.networks.nets import UNet # module is monai-specific; safe to import eagerly from src.common import dice # noqa: E402 # Selection gates / consistency filter shared with Pipeline B (pure numpy/scipy) from nnu_common import pos_mask, neg_frac_bg, consistency_filter, load_voxel_stats # noqa: E402,F401 WIN = 96 # ---------------- network / loss ---------------- def build_model(device, norm="batch", b_aligned=False): from monai.networks.nets import UNet if b_aligned: return UNetBAligned(device, norm=norm).to(device) net = UNet( spatial_dims=3, in_channels=1, out_channels=2, channels=(16, 32, 64, 128), strides=(2, 2, 2), num_res_units=2, act=("LEAKYRELU", {"inplace": True, "negative_slope": 0.01}), norm=("instance", {"affine": True, "eps": 1e-5}) if norm == "instance" else "batch", ) return net.to(device) class UNetBAligned(UNet): """Pipeline C's MONAI UNet + Pipeline B's deep supervision. Adds 1x1 conv aux heads at the 1/2- and 1/4-resolution decoder outputs (the 1/8 bottleneck is dropped, mirroring nnU-Net's zero-weight lowest scale). In training mode forward returns (final, aux_half, aux_quarter); in eval mode it returns just the final logits so MONAI sliding-window inference works unchanged. """ def __init__(self, device, norm="instance"): super().__init__( spatial_dims=3, in_channels=1, out_channels=2, channels=(16, 32, 64, 128), strides=(2, 2, 2), num_res_units=2, act=("LEAKYRELU", {"inplace": True, "negative_slope": 0.01}), norm=("instance", {"affine": True, "eps": 1e-5}) if norm == "instance" else "batch", ) # locate the up-paths: model[1].submodule = [down2, SkipConn(S3), up2] s2 = self.model[1].submodule self._up_half = s2[2] # 16ch @ 1/2 s3 = s2[1].submodule # [down3, SkipConn(bottom), up3] self._up_quarter = s3[2] # 32ch @ 1/4 self.aux_half = nn.Conv3d(16, 2, 1) self.aux_quarter = nn.Conv3d(32, 2, 1) self._aux = {} self._up_half.register_forward_hook(lambda m, i, o: self._aux.__setitem__("half", o)) self._up_quarter.register_forward_hook(lambda m, i, o: self._aux.__setitem__("quarter", o)) from src.unet3d import init_b self.apply(init_b) def forward(self, x): self._aux = {} out = super().forward(x) if not self.training: return out return out, self.aux_half(self._aux["half"]), self.aux_quarter(self._aux["quarter"]) class WeightedDiceCELoss(nn.Module): """Per-sample (soft-dice + CE), weighted batch mean. Same convention as Pipeline A's src/losses.py per_sample_loss (dice over background+foreground averaged, dice_weight 0.5), built on the MONAI DiceLoss. Per-sample reduction is required so accepted pseudo-label rows can be down-weighted individually via row "w"; a batch-level loss (MONAI DiceCELoss / nnU-Net CE+Dice) cannot express that. """ def __init__(self, dice_weight=0.5): super().__init__() from monai.losses import DiceLoss self.dice_weight = dice_weight self.dice = DiceLoss(include_background=True, softmax=False, reduction="none") def forward(self, logits, label, w): logits = logits.float() c = logits.size(1) one = nn.functional.one_hot(label, c).permute(0, -1, *range(1, label.ndim)).float() p = nn.functional.softmax(logits, dim=1) dsc = self.dice(p, one) # (B, C, 1, 1, 1) dsc = dsc.squeeze(-1).squeeze(-1).squeeze(-1).mean(dim=1) # (B,) ce = nn.functional.cross_entropy(logits, label, reduction="none") # (B, z, y, x) ce = ce.mean(dim=tuple(range(1, ce.ndim))) # (B,) per = (1 - self.dice_weight) * ce + self.dice_weight * dsc if torch.isnan(per).any() or torch.isinf(per).any(): per = torch.zeros_like(per) # zero-gradient fallback: keeps DDP collectives in sync return (per * w).sum() / w.sum().clamp(min=1e-6) # ---------------- data ---------------- class NNDCeLoss(nn.Module): """nnU-Net-style DC_and_CE (same formulation as src.losses.nn_dc_and_ce): unweighted CE + batch-level soft dice over both classes, smooth=1.0, loss = CE - mean(dice). Used for the A/C-vs-B loss-parity runs.""" def forward(self, logits, label): logits = logits.float() ce = nn.functional.cross_entropy(logits, label) p = nn.functional.softmax(logits, dim=1) g = nn.functional.one_hot(label, p.size(1)).permute(0, -1, *range(1, label.ndim)).float() axes = (0, *range(2, p.dim())) tp = (p * g).sum(axes) fp = (p * (1.0 - g)).sum(axes) fn = ((1.0 - p) * g).sum(axes) dc = (2.0 * tp + 1.0) / (2.0 * tp + fp + fn + 1.0) loss = ce - dc.mean() if not torch.isfinite(loss): return torch.zeros_like(loss) # zero-gradient fallback (DDP parity) return loss class _WorkerRNG: """Per-dataloader-worker numpy RandomState (forked workers reseed independently).""" @staticmethod def get(seed=0): info = torch.utils.data.get_worker_info() wid = info.id if info is not None else 0 return np.random.RandomState(seed + 104729 * wid) class ZScoreAll: """Per-volume Z-score over all voxels (Pipeline B's ZScoreNormalization, no mask). Applied to channel-first pimg.""" def __call__(self, data): from src.augment_b import zscore_volume d = dict(data) img = d["pimg"] v = img[0] if img.ndim == 4 else img v = zscore_volume(np.asarray(v, dtype=np.float32)) d["pimg"] = v[np.newaxis] if img.ndim == 4 else v return d class BAugmentB: """Pipeline B's full augmentation (src.augment_b), applied to the sampled 96^3 patch (B's SpatialTransform deforms the sampled patch about its center and crops back to the final size).""" def __call__(self, data): from src.augment_b import augment_b d = dict(data) img = np.asarray(d["pimg"], dtype=np.float32) raw_lab = d.get("plabel") lab = None if raw_lab is None else np.asarray(raw_lab) lab_4d = lab is not None and lab.ndim == 4 if lab_4d: lab = lab[0] out_img, out_lab = augment_b(img, lab, _WorkerRNG.get()) d["pimg"] = out_img if img.ndim == 4 else out_img[np.newaxis] if out_lab is not None: d["plabel"] = out_lab[np.newaxis] if lab_4d else out_lab return d def train_transform(win=WIN, b_aligned=False, no_aug=False): """Channel-first MONAI chain: load → min-pad → augment → pos/neg crop → typed tensors. Legacy (default): mirrors Pipeline A's augmentation set (per-axis flip p=0.5, rot90 in the (y,x) plane p=0.1, brightness 1±0.1 p=0.3, Gaussian noise σ=0.01 p=0.4). Patch sampling uses MONAI's foreground-aware RandCropByPosNegLabeld (pos=neg=1). b_aligned: Pipeline B's data protocol — per-volume Z-score, B's full augmentation (rot/zoom/noise/blur/brightness/contrast/lowres/gamma/mirror), and 33% foreground patch sampling (pos=neg ratio 0.5:1, like B's oversample_foreground_percent=0.33). """ from monai.transforms import ( Compose, LoadImaged, EnsureChannelFirstd, EnsureTyped, SpatialPadd, RandFlipd, RandRotate90d, RandScaleIntensityd, RandGaussianNoised, RandCropByPosNegLabeld, Compose as _C) if b_aligned: crop_pos, crop_neg = 0.5, 1.0 # B's SpatialTransform augments the SAMPLED patch (rot/zoom about its # center, then crop back to final size) — not the whole volume; that # is also what Pipeline A's b_aug does, and whole-volume deformation # is ~30x slower per item. steps = [ LoadImaged(keys=["pimg", "plabel"]), EnsureChannelFirstd(keys=["pimg", "plabel"]), SpatialPadd(keys=["pimg", "plabel"], spatial_size=(win, win, win)), ZScoreAll(), RandCropByPosNegLabeld(keys=["pimg", "plabel"], label_key="plabel", spatial_size=(win, win, win), pos=crop_pos, neg=crop_neg, num_samples=1), ] if not no_aug: steps.append(BAugmentB()) steps.append(EnsureTyped(keys=["pimg", "plabel"], dtype=[torch.float32, torch.long])) return Compose(steps) return Compose([ LoadImaged(keys=["pimg", "plabel"]), EnsureChannelFirstd(keys=["pimg", "plabel"]), SpatialPadd(keys=["pimg", "plabel"], spatial_size=(win, win, win)), RandFlipd(keys=["pimg", "plabel"], prob=0.5, spatial_axis=0), RandFlipd(keys=["pimg", "plabel"], prob=0.5, spatial_axis=1), RandFlipd(keys=["pimg", "plabel"], prob=0.5, spatial_axis=2), RandRotate90d(keys=["pimg", "plabel"], prob=0.1, max_k=3, spatial_axes=(1, 2)), RandScaleIntensityd(keys=["pimg"], factors=0.1, prob=0.3), RandGaussianNoised(keys=["pimg"], std=0.01, prob=0.4), RandCropByPosNegLabeld(keys=["pimg", "plabel"], label_key="plabel", spatial_size=(win, win, win), pos=1, neg=1, num_samples=1), EnsureTyped(keys=["pimg", "plabel"], dtype=[torch.float32, torch.long]), ]) def rows_for_dataset(rows): out = [] for r in rows: out.append({"key": r["key"], "pimg": r["pimg"], "plabel": r["plabel"], "w": float(r.get("w", 1.0))}) return out def _collate(batch): # RandCropByPosNegLabeld(num_samples=1) yields a 1-item list per sample items = [b[0] if isinstance(b, (list, tuple)) else b for b in batch] return torch.utils.data.default_collate(items) def make_dataloader(rows, win=WIN, batch=3, workers=4, b_aligned=False): from monai.data import Dataset ds = Dataset(data=rows_for_dataset(rows), transform=train_transform(win, b_aligned=b_aligned)) return torch.utils.data.DataLoader( ds, batch_size=batch, shuffle=True, num_workers=workers, collate_fn=_collate, drop_last=len(ds) > batch, pin_memory=True, persistent_workers=workers > 0) # ---------------- inference ---------------- @torch.no_grad() def predict_probs(model, vol, device, win=WIN, overlap=0.5, tta=True, sw_batch=8, zscore=False): """Tumor probability map (C, z, y, x → (z, y, x) numpy) for one volume. MONAI sliding_window_inference (gaussian blend, `overlap`), plus the same 4-view flip TTA (identity + 3 axis flips) as Pipeline A. zscore=True matches a B-aligned model's input (per-volume Z-score, no [0,1.5] clip). """ from monai.inferers import sliding_window_inference a = np.nan_to_num(np.asarray(vol, dtype=np.float32), nan=0.0, posinf=1.5, neginf=0.0) if zscore: from src.augment_b import zscore_volume a = zscore_volume(a) else: a = np.clip(a, 0.0, 1.5) t = torch.from_numpy(a).unsqueeze(0).unsqueeze(0).to(device) def run(vt): with torch.autocast("cuda", dtype=torch.bfloat16): logits = sliding_window_inference( vt, roi_size=(win, win, win), sw_batch_size=sw_batch, predictor=model, mode="gaussian", overlap=overlap, sw_device=device, device=device) return torch.softmax(logits.float(), dim=1)[0, 1].cpu().numpy() acc = run(t) if not tta: return acc for ax in (0, 1, 2): acc += np.flip(run(t.flip(ax + 2)), axis=ax) return acc / 4.0 def prob_dice(probs, lab, thr=0.5): """Dice at probability threshold 0.5 — the convention of Pipeline A's 08_eval and Pipeline B's primary metric.""" return dice((probs >= thr).astype("uint8"), (lab > 0).astype("uint8")) def ckpt_b_aligned(sd): """Whether a checkpoint was trained with the B-aligned config (instance norm + deep-sup aux heads). New checkpoints carry an explicit 'b_aligned' marker; pre-marker ones fall back to the aux_head state-dict keys.""" if not isinstance(sd, dict): return False if "b_aligned" in sd: return bool(sd["b_aligned"]) inner = sd.get("model", sd) return isinstance(inner, dict) and any(k.startswith("aux_") for k in inner) def load_model(ckpt, device): sd = torch.load(ckpt, map_location=device, weights_only=True) ba = ckpt_b_aligned(sd) net = build_model(device, norm="instance" if ba else "batch", b_aligned=ba) net.load_state_dict(sd.get("model", sd)) net.eval() return net, sd # ---------------- distributed (torchrun, as in Pipeline A) ---------------- def rank_info(): return (int(os.environ.get("RANK", 0)), int(os.environ.get("WORLD_SIZE", 1)), int(os.environ.get("LOCAL_RANK", 0))) def init_dist(): if int(os.environ.get("WORLD_SIZE", 1)) > 1: from datetime import timedelta import torch.distributed as dist dist.init_process_group("nccl", timeout=timedelta(minutes=30)) def barrier(): if int(os.environ.get("WORLD_SIZE", 1)) > 1: import torch.distributed as dist dist.barrier() def destroy_dist(): if int(os.environ.get("WORLD_SIZE", 1)) > 1: import torch.distributed as dist dist.destroy_process_group()