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