| --- |
| license: apache-2.0 |
| tags: |
| - medical-imaging |
| - segmentation |
| - image-segmentation |
| - ct |
| - lung |
| - nodule |
| - monai |
| - pytorch |
| library_name: pytorch |
| pipeline_tag: image-segmentation |
| --- |
| |
| # 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](https://docs.monai.io/en/stable/networks.html#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. |
|
|
| | Metric | Value | |
| |---------------------|---------:| |
| | mIoU | 0.7509 | |
| | Accuracy | 0.9995 | |
| | Precision | 0.6077 | |
| | Recall | 0.7434 | |
| | Dice / F1 (micro) | 0.6687 | |
| | IoU (nodule class) | 0.5023 | |
| | Dice per case, mean | 0.5250 | |
| | Dice per case, med. | 0.5853 | |
|
|
| 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 |
|
|
| ```python |
| 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 (T_max = 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: |
| |
| ```bash |
| 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. |
| |
