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

Joint Lung + Nodule Segmentation — SegResNet (3D, ex-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 NLST + NSCLC-Radiomics subset of the unified corpus — LIDC-IDRI is excluded because it lacks ground-truth 2D lung labels.

Companion to joint-segresnet-3d-pseudo-lidc (same architecture, LIDC included via ROI-model-derived pseudo-lung labels), 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_ex_lidc split (patient-grouped, dataset-stratified, LIDC excluded):

Source	Role
NLST	train + val
NSCLC-Radiomics	train + val
LIDC-IDRI	excluded

Split sizes: 1 110 train / 196 val / 129 test (held out).

Lung labels are ground-truth per-slice 2D masks (from roi_sem_seg_2d/), stacked into 3D per series. Nodule labels come from the corpus's own 3D nodule annotations.

Validation metrics (best-epoch, val split)

Class	Dice	Recall	Precision
Lung	0.9814	0.991	0.973
Nodule	0.6546	0.650	0.659
Combined (mean)	0.8180	—	—

Combined score = 0.5 · (Lung Dice + Nodule Dice). Both class Dices are per-case-averaged on the val split. Full metric breakdown (mIoU, Accuracy, per-class Precision/Recall, per-case distribution) is included as metrics.json alongside this model card once the eval pass completes.

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: ≈ 3.5 days

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

Ablation: ex-LIDC vs pseudo-LIDC

The joint task cannot use LIDC directly because LIDC lacks GT lung labels. Two variants were trained:

• This model (ex-LIDC): drop LIDC entirely → 1 110 train cases, all with full GT.
• joint-segresnet-3d-pseudo-lidc: include LIDC with lung labels produced by the 2D SegResNet ROI model → 1 683 train cases, mixed GT
  • pseudo.

On the val split, ex-LIDC beats pseudo-LIDC by 0.009 combined Dice (0.8180 vs 0.8087). The pseudo-labels' noise slightly hurts the lung head's supervision signal; the additional LIDC diversity does not compensate. See the companion model card for details.

License & intended use

Model weights released under Apache 2.0. Training data was public but covered by dataset-specific terms (NLST, NSCLC-Radiomics) — 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.