"""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 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): from monai.networks.nets import UNet 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="batch", ) return net.to(device) 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 ---------------- def train_transform(win=WIN): """Channel-first MONAI chain: load → min-pad → augment → pos/neg crop → typed tensors. 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) instead of Pipeline A's uniform random crop. """ from monai.transforms import ( Compose, LoadImaged, EnsureChannelFirstd, EnsureTyped, SpatialPadd, RandFlipd, RandRotate90d, RandScaleIntensityd, RandGaussianNoised, RandCropByPosNegLabeld) 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): from monai.data import Dataset ds = Dataset(data=rows_for_dataset(rows), transform=train_transform(win)) 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): """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. """ 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) 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 load_model(ckpt, device): sd = torch.load(ckpt, map_location=device, weights_only=True) net = build_model(device) 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()