HalmosiL/joint-segresnet-3d-pseudo-lidc

Joint Lung + Nodule Segmentation — SegResNet (3D, pseudo-LIDC)

Single-model, end-to-end 3-class segmentation of chest CT: predicts background / lung / nodule in one forward pass, directly from the full CT (no ROI stage, no bbox crop). Trained on the full unified NLST + NSCLC + LIDC-IDRI corpus; because LIDC lacks ground-truth 2D lung labels, LIDC lung supervision is supplied as pseudo-GT predictions from the 2D SegResNet ROI model.

Companion to joint-segresnet-3d-ex-lidc (same architecture, LIDC dropped entirely instead of pseudo-labeled), and to joint-dynunet-3d-{ex,pseudo}-lidc (same task, different architecture).

Model details

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

Nodule is treated as a class distinct from lung: a voxel that is both lung tissue and nodule is assigned exclusively to class 2 (nodule takes precedence over lung).

Data

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

Source	Role	Lung labels
NLST	train + val	GT (per-slice 2D masks stacked to 3D)
NSCLC-Radiomics	train + val	GT (per-slice 2D masks stacked to 3D)
LIDC-IDRI	train + val	pseudo-GT — produced by Kakimaki00/roi-segresnet-2d on each LIDC slice, then stacked

Split sizes: 1 683 train / 297 val / 325 test (held out).

Nodule labels come from the corpus's own 3D nodule annotations for all three sources.

Validation metrics (best-epoch, val split)

Class	Dice	Recall	Precision
Lung	0.9754	0.983	0.967
Nodule	0.6421	0.643	0.641
Combined (mean)	0.8087	—	—

Combined score = 0.5 · (Lung Dice + Nodule Dice). Both class Dices are per-case-averaged on the val split.

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"],   # 3
    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, 256)             # (B, C, H, W, D)
    logits = model(x)                                # (B, 3, D, H, W)
    pred_class = logits.argmax(dim=1)                # (B, D, H, W) in {0, 1, 2}
    lung_mask   = (pred_class == 1).to(torch.uint8)
    nodule_mask = (pred_class == 2).to(torch.uint8)

Unlike the two-stage nodule-* checkpoints in this collection, this model does not need a lung-bbox crop — feed it the whole CT resampled to 256³.

Training recipe

• Loss: Multi-class Focal Tversky + weighted CE (α=0.3, β=0.7, γ=2.0, λ_CE=0.1, class weights = [1.0, 1.0, 100.0] for [bg, lung, nodule])
• Optimizer: Adam (lr = 1e-5, weight decay = 1e-5)
• Scheduler: CosineAnnealingLR (T_max = 400, η_min = 1e-6)
• Batch size: 4
• Epochs: 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: ≈ 5 days (more train cases than the ex-LIDC variant)

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

Ablation: ex-LIDC vs pseudo-LIDC

Compared to joint-segresnet-3d-ex-lidc (same architecture, LIDC dropped): the pseudo-LIDC variant is 0.009 combined Dice worse (0.8087 vs 0.8180) despite training on 51 % more cases. The pseudo-labels' noise slightly hurts the lung head's supervision signal; the additional LIDC diversity does not compensate. If you need a joint model, the ex-LIDC variant is the recommended default.

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.