Lung ROI Segmentation β€” SegResNet (2D)

Per-slice 2D lung foreground (thoracic parenchyma) segmentation for chest CT. Used as the ROI stage of a two-stage coarse-to-fine pulmonary-nodule segmentation pipeline: its outputs are stacked into a 3D lung bounding box that the downstream nodule model crops to.

Model details

  • Architecture: MONAI SegResNet, 2D residual U-Net
  • Trainable parameters: 6,904,081
  • 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_v2 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.

Validation metrics

Evaluated on 64 097 val slices (β‰ˆ 4.2 Γ— 10⁹ pixels), micro-averaged at 0.5 threshold on sigmoid output.

Metric Value
mIoU 0.9803
Accuracy 0.9953
Precision 0.9797
Recall 0.9857
Dice (F1) 0.9827

How to load & run inference

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

cfg = yaml.safe_load(open("config.yaml"))["model"]
model = SegResNet(
    spatial_dims = cfg["spatial_dims"],
    in_channels  = cfg["in_channels"],
    out_channels = cfg["out_channels"],
    init_filters = cfg["init_filters"],
    blocks_down  = tuple(cfg["blocks_down"]),
    blocks_up    = tuple(cfg["blocks_up"]),
    dropout_prob = cfg["dropout_prob"],
)
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-3, weight decay = 1e-5)
  • Scheduler: CosineAnnealingLR (T_max = 100, Ξ·_min = 1e-6)
  • Batch size: 16
  • Samples/epoch: 20 000 (random subsample of ~365 k train slices)
  • Epochs budget: 100 | best checkpoint at epoch 7 / 100
  • 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.

Downloads last month
10
Inference Providers NEW
This model isn't deployed by any Inference Provider. πŸ™‹ Ask for provider support