3D CT Segmentation Pipeline

This repository provides a full 3D CT segmentation workflow built with PyTorch and MONAI:

• 3D U-Net-style model reconstructed from the provided architecture image
• training with mixed precision and resume support
• validation Dice tracking and CSV logging
• TensorBoard logging
• inference for .nii.gz volumes
• Gradio demo for interactive prediction
• optional SLURM launch script
• optional Hugging Face export

Data layout

imagesTr/Patient_01_0000.nii.gz
imagesTr/Patient_02_0000.nii.gz
labelsTr/Patient_01.nii.gz
labelsTr/Patient_02.nii.gz

Every image volume must have a matching label volume.

Installation

conda create -n ctseg python=3.10 -y
conda activate ctseg
pip install -r requirements.txt

Training

Smoke test

python train.py \
  --images-tr /data/imagesTr \
  --labels-tr /data/labelsTr \
  --outdir outputs \
  --num-classes 7 \
  --epochs 1 \
  --batch-size 1 \
  --amp \
  --num-workers 0 \
  --cache-rate 0.0

Full run

python train.py \
  --images-tr /data/imagesTr \
  --labels-tr /data/labelsTr \
  --outdir outputs \
  --num-classes 7 \
  --epochs 100 \
  --batch-size 1 \
  --amp \
  --num-workers 4 \
  --cache-rate 0.0

Resume training

python train.py \
  --images-tr /data/imagesTr \
  --labels-tr /data/labelsTr \
  --outdir outputs \
  --num-classes 7 \
  --epochs 200 \
  --batch-size 1 \
  --amp \
  --resume outputs/checkpoints/best.pt

Outputs

The training script writes:

• outputs/metrics.csv — epoch, training loss, validation Dice
• outputs/tb_logs/ — TensorBoard event files
• outputs/checkpoints/last.pt — latest checkpoint
• outputs/checkpoints/best.pt — best checkpoint by validation Dice
• outputs/best.pt — convenience copy of the best checkpoint
• outputs/hparams.json — run configuration

TensorBoard

tensorboard --logdir outputs/tb_logs --port 6006

Inference

python infer.py \
  --input /data/Patient_99_0000.nii.gz \
  --checkpoint outputs/best.pt \
  --output /data/Patient_99_pred.nii.gz \
  --num-classes 7 \
  --in-channels 1

Gradio demo

Set the checkpoint path and launch:

export CHECKPOINT_PATH=weights/best.pt
export NUM_CLASSES=7
python app.py

Export for Hugging Face

python export_hf.py --checkpoint outputs/best.pt --outdir hf_model
python export_hf.py --checkpoint outputs/best.pt --outdir hf_model --push --repo-id your-name/ct-seg-model

Notes

• --num-classes must match your label values. For labels with values 0..6, use --num-classes 7.
• Use --batch-size 1 for 3D volumes unless you have tested larger batches.
• The model uses fixed-size preprocessing with orientation normalization and padding/cropping to keep batching stable.