szabopeter/nodule-segresnet-3d-small

Lung Nodule Segmentation — SegResNet (3D, small)

Voxel-level segmentation of pulmonary nodules in 3D chest CT volumes, cropped to the lung region. Small-capacity SegResNet (init_filters = 16, ~20 M params) trained on lung-bbox crops from the unified NLST + NSCLC + LIDC-IDRI corpus.

Named segresnet_small in the demo repo.

Model details

• —Architecture: MONAI SegResNet, 3D residual U-Net
• —Trainable parameters: 20,663,538
• —Input: (1, 256, 256, 256) CT crop, intensity-normalised to [0, 1], resampled from a per-series lung bbox (padding = 20 vox)
• —Output: (2, 256, 256, 256) softmax logits — class 0 = background, class 1 = nodule
• —Framework: PyTorch + MONAI

Data

Trained on the unified split (patient-grouped, dataset-stratified):

• —NLST
• —NSCLC-Radiomics
• —LIDC-IDRI

Split sizes: 1 683 train / 297 val / 325 test (held out). Nodule voxels are on the order of ~10⁻⁵ of the total — see the note on metric interpretation below.

Validation metrics

Evaluated on 297 val volumes (≈ 5 × 10⁹ voxels), micro-averaged at the argmax of the 2-class softmax output.

Because ~10⁻⁵ of voxels are nodule, Accuracy is trivially near 1.0 regardless of model quality and mIoU is dominated by IoU_background. The operationally meaningful numbers are Precision, Recall, and the per-case Dice distribution.

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():
    # Input: (B, 1, 256, 256, 256), pre-cropped to the lung bbox and normalised to [0, 1]
    x = torch.randn(1, 1, 256, 256, 256)
    logits = model(x)                          # (B, 2, D, H, W)
    pred_class = logits.argmax(dim=1)          # (B, D, H, W)
    nodule_mask = (pred_class == 1).to(torch.uint8)

The model expects lung-bbox-cropped input. A separate 2D ROI model is needed to produce that bbox — see the accompanying ROI checkpoints (roi-segresnet-2d or roi-swinunetr-2d).

Training recipe

• —Loss: Focal Tversky + weighted CE (α=0.3, β=0.7, γ=2.0, λ_CE=0.1, nodule class weight = 100)
• —Optimizer: Adam (lr = 1e-5, weight decay = 1e-5)
• —Scheduler: CosineAnnealingLR (Tmax = 400, ηmin = 1e-6)
• —Batch size: 4
• —Epochs: 400 | best checkpoint at epoch 339 / 400
• —Mixed precision: bf16
• —Augmentation: 3D flips, 90° rotations, elastic rotation, zoom, intensity scale/shift, Gaussian noise/blur, contrast
• —Seed: 42
• —Hardware: 1 × NVIDIA H100 94 GB
• —Wall-clock: ≈ 3.5 days

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_3d/ and nodule_sem_seg_3d/
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.