roi-swinunetr-2d

Lung ROI Segmentation — SwinUNETR (2D, small)

Per-slice 2D lung foreground segmentation for chest CT — Swin-transformer variant. Same task, same data, same eval as the SegResNet ROI model, but with a Swin encoder

• CNN decoder instead of a plain residual CNN. Small-capacity variant (feature_size = 24) chosen to match the SegResNet baseline's parameter budget for an apples-to-apples architecture comparison.

Model details

• Architecture: MONAI SwinUNETR, Swin transformer + CNN decoder
• Trainable parameters: 6,302,203
• Input: (1, 256, 256) axial CT slice, intensity-normalised to [0, 1]
• Output: (1, 256, 256) sigmoid map; foreground = lung tissue
• Framework: PyTorch + MONAI

Data

Trained on the unified split (patient-grouped, dataset-stratified) of a unified corpus assembled from three public sources:

• NLST
• NSCLC-Radiomics
• LIDC-IDRI

Slices per split: 365 014 train / 64 097 val / — test held out. Identical split and sampling recipe as the SegResNet baseline; only the model family and LR differ.

Validation metrics

Evaluated on 64 097 val slices (≈ 4.2 × 10⁹ pixels), micro-averaged at 0.5 threshold on sigmoid output.

The largest improvement over the SegResNet baseline is in the tail of the per-slice distribution (Dice_p05 rises from 0.00 to 0.88) — top-/bottom-of-volume slices where the lung is nearly absent are handled more conservatively.

How to load & run inference

import yaml, torch
from monai.networks.nets import SwinUNETR

cfg = yaml.safe_load(open("config.yaml"))["model"]
model = SwinUNETR(
    in_channels    = cfg["in_channels"],
    out_channels   = cfg["out_channels"],
    feature_size   = cfg["feature_size"],
    depths         = tuple(cfg["depths"]),
    num_heads      = tuple(cfg["num_heads"]),
    spatial_dims   = cfg["spatial_dims"],
    use_checkpoint = cfg.get("use_checkpoint", False),
)
state = torch.load("model.pth", map_location="cpu", weights_only=True)
model.load_state_dict(state)
model.eval()

with torch.no_grad():
    x = torch.randn(1, 1, 256, 256)         # (B, C, H, W) — replace with your CT slice
    prob = torch.sigmoid(model(x))
    lung_mask = (prob > 0.5).to(torch.uint8)

Training recipe

• Loss: DiceLoss(sigmoid=True, squared_pred=True)
• Optimizer: Adam (lr = 1e-4, weight decay = 1e-5) — transformer families are more LR-sensitive than CNNs
• Scheduler: CosineAnnealingLR (T_max = 100, η_min = 1e-6)
• Batch size: 16
• Samples/epoch: 20 000
• Epochs budget: 100 | best checkpoint at epoch 59 / 100 (training was still in progress at eval time; val Dice had plateaued in [0.9647, 0.9655] for ~15 preceding epochs)
• Mixed precision: bf16
• Seed: 42
• Hardware: 1 × NVIDIA H100 94 GB

Full config is included in this repo as config.yaml.

Reproducing training

Training code lives in an accompanying reproduction demo (published separately). Once available, reproduce with:

export DATA_ROOT=/path/to/unified          # dir containing ct_2d/ and roi_sem_seg_2d/
python train.py --config config.yaml

License & intended use

Model weights released under Apache 2.0. Training data was public but covered by dataset-specific terms (NLST, NSCLC-Radiomics, LIDC-IDRI) — users must comply with those separately when using the model on comparable data.

Not a medical device. Not intended for clinical use. Research only.

Citation

Paper in preparation.