Skip to content

Pretrained Model Weights

Latest
Latest

Choose a tag to compare

View all tags
@Albertzry Albertzry released this 07 Mar 02:54
· 2 commits to ldh-alternative since this release
weights
1961afe

Release: r20250307

This release provides all pretrained model weights for TotalSpineSeg-v2, which supports full-spine MRI segmentation and lumbar intervertebral disc (IVD) degeneration detection & segmentation.

📦 Download Files

File Size Description
Dataset101_TotalSpineSeg_step1.tar.gz ~221 MB Step 1 — Coarse Segmentation & Landmarks. nnU-Net model (9-class segmentation: cord, canal, IVDs, vertebrae, + 5 landmark classes). Trainer: nnUNetTrainer_DASegOrd0_NoMirroring, Plans: nnUNetPlans_small, Config: 3d_fullres, Fold 0.
Dataset102_TotalSpineSeg_step2.tar.gz ~222 MB Step 2 — Fine-Grained Labeling. nnU-Net model (11-class segmentation with odd/even vertebrae, sacrum, and 4 landmark discs). Uses Step 1 output as second input channel. Same trainer/plans as Step 1, Fold 0.
Dataset105_TotalSpineSeg_LDH.tar.gz ~7.5 MB Stage 3A — IVD Degeneration Detection. Custom StageADetectorV2 (Residual 3D CNN + SE attention, Focal Loss). Performs disc-level binary classification for IVD degeneration. Fold 0.
Dataset107_LDH_ROI.tar.gz ~231 MB Stage 3B — IVD Degeneration Fine Segmentation. nnU-Net model (3d_fullres) for ROI-based IVD degeneration mask prediction. 3-channel input (MRI + disc mask + disc index map). Includes both checkpoint_best.pth and checkpoint_final.pth, Fold 0.

🛠️ Installation

  1. Download all 4 .tar.gz files from this release.

  2. Set your data directory:

    export TOTALSPINESEG_DATA="/path/to/your/data"

  3. Extract into the nnU-Net results folder:

    mkdir -p $TOTALSPINESEG_DATA/nnUNet/results
cd $TOTALSPINESEG_DATA/nnUNet/results

tar xzf /path/to/Dataset101_TotalSpineSeg_step1.tar.gz
tar xzf /path/to/Dataset102_TotalSpineSeg_step2.tar.gz
tar xzf /path/to/Dataset105_TotalSpineSeg_LDH.tar.gz
tar xzf /path/to/Dataset107_LDH_ROI.tar.gz

  4. Verify the expected directory structure:

    $TOTALSPINESEG_DATA/nnUNet/results/
├── Dataset101_TotalSpineSeg_step1/
│   └── nnUNetTrainer_DASegOrd0_NoMirroring__nnUNetPlans_small__3d_fullres/
│       └── fold_0/checkpoint_final.pth
├── Dataset102_TotalSpineSeg_step2/
│   └── nnUNetTrainer_DASegOrd0_NoMirroring__nnUNetPlans_small__3d_fullres/
│       └── fold_0/checkpoint_final.pth
├── Dataset105_TotalSpineSeg_LDH/
│   └── ldh_twostage/checkpoints/ldh_stageA_fold_0.pth
└── Dataset107_LDH_ROI/
    └── nnUNetTrainer__nnUNetPlans__3d_fullres/
        └── fold_0/{checkpoint_best.pth, checkpoint_final.pth}

🚀 Quick Start

# Full spine segmentation (Step 1 + Step 2)
totalspineseg --input-dir /path/to/mri_images --output-dir /path/to/output

# IVD degeneration detection & segmentation (Stage 3A + 3B)
python scripts/infer_ldh.py --input-dir /path/to/mri_images --output-dir /path/to/output

# Clinical measurement report
python calculate.py --input-dir /path/to/output --output-dir /path/to/reports

📋 Model Architecture Summary

  • Step 1 & Step 2: nnU-Net v2 backbone with DASegOrd0_NoMirroring trainer variant and nnUNetPlans_small configuration, both using 3d_fullres.
  • Stage 3A: Custom StageADetectorV2 — a residual 3D CNN with Squeeze-and-Excitation (SE) attention blocks, trained with Focal Loss (γ=2) for disc-level binary classification.
  • Stage 3B: Standard nnU-Net v2 (nnUNetTrainer, nnUNetPlans, 3d_fullres) trained on 3-channel 64³ ROI patches (MRI intensity + disc mask + disc index map) for voxel-level IVD degeneration segmentation.

⚠️ Requirements

  • Python ≥ 3.10
  • PyTorch < 2.6 (CUDA recommended)
  • nnU-Net v2 (pip install -e '.[nnunetv2]')

See README.md for full installation instructions.

Assets 6
Loading