ribline-3d-unet: Lightweight 3D U-Net for Thoracic CT Rib Centerline Probability Estimation

Official Lightweight 3D U-Net checkpoint (rib_unet_3d.pth) and 24-rib statistical atlas for ribline (v0.3.0), described in Hybrid Context Forests and Bilateral Viterbi Matching for 3D Rib Centerline Tracing in CT (PDF):

Hugging Face Repository	Model Variant	12-Scan LCD (mm)	12-Scan Acc (%)	5-Fold CV LCD (mm)	5-Fold CV Acc (%)
wq2012/ribline-3d-ct-centerline	Hybrid SCF + 3D U-Net + Stage 3 (Full Suite)	2.05 +/- 0.48	97.2 %	3.64 +/- 3.90	91.8 %
wq2012/ribline-anatomical-scf	Anatomically Anchored SCF + Stage 3 (CPU-only)	2.88 +/- 1.32	95.1 %	5.28 +/- 4.62	86.9 %
wq2012/ribline-3d-unet	Lightweight 3D U-Net + Stage 3	3.40 +/- 1.99	92.0 %	7.44 +/- 8.93	78.2 %
wq2012/ribline-scf-pass2-autocontext	2-Pass Auto-Context SCF Baseline + Stage 3	9.36 +/- 10.57	71.0 %	--	--

Architecture and Training Recipe

• Backbone (RibUNet / RibUNet3D): Plain 3-level 3D encoder-decoder (16 -> 32 -> 64 channels, two 3x3x3 Conv3d -> GroupNorm(min(8, C)) -> ReLU layers per stage, 2x2x2 MaxPool3d downsampling, 2x2x2 ConvTranspose3d upsampling with skip concatenations, and a 1x1x1 output head; 339,889 trainable parameters, no residual blocks).
• Training (train_config_unet.json): Trained on 64x64x64 patches (70 % foreground-centered, 30 % uniform random) with per-volume z-score normalization of raw stored intensities, optimizing 0.5 * SoftDice + 0.5 * BCEWithLogits with Adam (lr = 1e-3, batch size 2, 8 patches/scan/epoch, 20 epochs, no data augmentation) against a 2x morphological dilation of the 1-voxel ground-truth centerlines.
• Provenance: rib_unet_3d.pth (md5 9e456ba9ff81c402a58436c666f152ef) is the exact epoch-20 checkpoint that produced the released validation probability maps (max_abs_diff = 0.0 upon re-inference).

Included Files

File	Description
rib_unet_3d.pth	Trained 3-level 3D U-Net PyTorch state dict (339,889 parameters).
train_config_unet.json	Training recipe configuration for rib_unet_3d.pth.
rib_atlas.joblib	24-rib statistical shape & topological atlas fitted on the 47 training scans.
benchmark_val12_fixed_gt.json, benchmark_val12_original_gt.json, benchmark_gt_original_excl_rib007_R07.json, benchmark_noSko.json, benchmark_train47.json, benchmark_cv5.json	Evaluation reports across the 12-scan eval split, 47-scan train split, and 59-scan 5-fold CV.
polar_stats_all59.json, gt_fix_log.json	Azimuthal monotonicity statistics and log of the single ground-truth fix.

Quick Start

import torch
from ribline import load_pretrained_models

# Downloads and instantiates the 24-rib atlas and 3D U-Net
models = load_pretrained_models("wq2012/ribline-3d-unet", load_scf=False, device="cpu")
unet = models["unet"]

# Predict probabilities in [0, 1] on a (B, 1, D, H, W) tensor
x = torch.randn(1, 1, 64, 64, 64)
probs = unet.predict_proba(x)

Or run CLI sliding-window inference directly:

python -m ribline.scripts.infer_unet \
    --model pretrained_models/rib_unet_3d.pth \
    --image-dir data/vol2mm --split split.json --out-dir predictions/unet

Evaluation Protocol and Results

All models were trained on the 47 annotated training scans (1,116 ribs) of the Siemens Corporate Research (SCR) 2 mm thoracic CT benchmark (split.json) and evaluated on the 12 evaluation scans (286 ribs after the documented ground-truth fix in RIB_007, plus 5-fold cross-validation across all 59 annotated scans / 1,402 ribs).

• Labeled Chamfer distance (LCD): mean over all GT ribs of 0.5 * (d_p2g + d_g2p) between each GT rib and the predicted curve with the same label (R01..R12, L01..L12); a GT rib without a predicted curve of that label receives a 50 mm penalty.
• Symmetric 10 mm labeling accuracy: a GT rib is counted as correctly labeled if and only if a same-label prediction exists, d_p2g <= 10 mm and d_g2p <= 10 mm, and the GT rib is the closest GT rib to that prediction.

12-Scan Evaluation Ablation (benchmark_val12_fixed_gt.json, 286 GT ribs)

Method	LCD (mm)	95 % CI	UCD (mm)	lab. clDice 3 / 4 mm	Acc all	Acc 1st	Acc 2-11	Acc 12th	Recall	Precision (FP)	NLD
RF Pass-1 + Z-sort	14.10 +/- 11.10	[7.96, 20.46]	3.31	0.548 / 0.606	64.3 %	79.2 %	62.9 %	63.6 %	91.3 %	95.5 % (13)	0.264
SCF Pass-2 + Z-sort	21.30 +/- 13.33	[13.79, 28.67]	4.16	0.416 / 0.452	43.0 %	45.8 %	44.2 %	27.3 %	76.2 %	100.0 % (0)	0.254
SCF Pass-2 + Stage 3	9.36 +/- 10.57	[4.24, 16.19]	3.73	0.683 / 0.720	71.0 %	54.2 %	73.8 %	59.1 %	78.0 %	99.6 % (1)	0.226
Anatomical SCF + Stage 3	2.88 +/- 1.32	[2.23, 3.73]	2.41	0.845 / 0.915	95.1 %	91.7 %	95.0 %	100.0 %	95.1 %	99.7 % (1)	0.088
3D U-Net + Z-sort	7.57 +/- 7.43	[3.97, 12.17]	2.75	0.710 / 0.813	83.6 %	70.8 %	86.7 %	63.6 %	92.0 %	100.0 % (0)	0.147
3D U-Net + Stage 3	3.40 +/- 1.99	[2.42, 4.64]	2.30	0.835 / 0.902	92.0 %	66.7 %	95.4 %	81.8 %	93.0 %	100.0 % (0)	0.121
Hybrid + Viterbi (no RibSnake)	2.85 +/- 1.70	[2.07, 3.95]	2.11	0.872 / 0.933	95.8 %	83.3 %	97.5 %	90.9 %	95.8 %	99.6 % (1)	0.147
Hybrid + full Stage 3	2.05 +/- 0.48	[1.80, 2.33]	1.91	0.886 / 0.945	97.2 % (278/286)	91.7 %	97.5 %	100.0 %	97.2 %	100.0 % (0)	0.068

5-Fold Cross-Validation Across All 59 Annotated Scans (benchmark_cv5.json, 1,402 GT ribs)

Method	LCD (mm)	95 % CI	Median LCD	UCD (mm)	lab. clDice 3 / 4 mm	Acc all	Acc 1st	Acc 2-11	Acc 12th	Recall	Precision (FP)	NLD
Anatomical SCF + Stage 3	5.28 +/- 4.62	[4.16, 6.56]	3.08	3.44	0.769 / 0.840	86.9 % (1218/1402)	72.6 %	87.7 %	93.0 %	90.8 %	95.3 % (67)	0.126
3D U-Net + Stage 3	7.44 +/- 8.93	[5.31, 9.85]	2.99	3.30	0.745 / 0.785	78.2 % (1097/1402)	61.5 %	80.9 %	67.8 %	87.0 %	98.1 % (26)	0.170
Hybrid + Viterbi (no RibSnake)	4.17 +/- 4.01	[3.29, 5.28]	2.78	3.07	0.807 / 0.879	91.4 % (1282/1402)	81.2 %	92.2 %	93.9 %	94.3 %	97.9 % (29)	0.195
Hybrid + full Stage 3	3.64 +/- 3.90	[2.79, 4.74]	2.39	2.85	0.830 / 0.889	91.8 % (1287/1402)	82.9 %	92.2 %	96.5 %	94.6 %	97.7 % (33)	0.097

Limitations and Protocol Caveats

• 60-scan single-source dataset at 2 mm isotropic spacing; no acquisition metadata; generalization outside SCR is untested.
• The 12-scan evaluation split is the alphabetical tail of the file list (non-random), and the Stage-3 heuristic constants were tuned while inspecting metrics on those 12 scans (while the forests, U-Net, and atlas were fitted strictly on the 47 training scans). See the 5-fold cross-validation table above for out-of-fold generalization across all 59 scans.
• Known annotation issues: merged R07+R08 in RIB_007 (fixed in gt_fix_log.json), inconsistent 12th ribs (absent in Tsc_004002), and possible patient overlap between Bro_004029 (train) and Sko_004029 (eval).

Citation

• Paper: https://arxiv.org/abs/2610.16047 (PDF)

@article{wang2026hybrid,
  title={Hybrid Context Forests and Bilateral Viterbi Matching for {3D} Rib Centerline Tracing in {CT}},
  author={Wang, Quan},
  journal={arXiv preprint arXiv:2610.16047},
  year={2026},
  url={https://arxiv.org/abs/2610.16047}
}