Pretrained Model Weights

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.

conda pack

🐍 Conda Environment Backup (Portable)

This release includes a fully packed Conda environment (tss) containing PyTorch, CUDA, nnU-Net, and all other necessary dependencies to run TotalSpineSeg-v2 out-of-the-box. Due to GitHub's 2GB file size limit, the 6.2GB archive has been split into 4 parts.

How to extract and use the environment:

1. Download all 4 parts (part_aa, part_ab, part_ac, part_ad) to the same directory.
2. Combine the split files back into the original archive:

   cat tss_env_packed.tar.gz.part_* > tss_env_packed.tar.gz