Completeness-Aware Reconstruction Enhancement

Diffusion based Anatomy-aware enhancement of sparse-view CT reconstruction.

The Completeness-Aware Reconstruction Enhancement (CARE) framework addresses a critical gap in sparse-view CT reconstruction by shifting the evaluation from traditional pixel-wise metrics to anatomy-aware metrics derived from automated structural segmentation. By incorporating segmentation-informed losses into latent diffusion models, CARE significantly improves the reconstruction fidelity of clinically relevant anatomical structures, ensuring that critical diagnostic features are preserved under highly limited view conditions.

Paper

Are Pixel-Wise Metrics Reliable for Sparse-View Computed Tomography Reconstruction?
Tianyu Lin¹, Xinran Li¹, Chuntung Zhuang¹, Qi Chen¹, Yuanhao Cai¹, Kai Ding², Alan L. Yuille¹ and Zongwei Zhou^1,*
1Johns Hopkins University, ²Johns Hopkins Medicine

We have documented common questions for the paper in Frequently Asked Questions (FAQ).

Installation

Create an conda environement via:

conda create -n care python=3.11 -y
conda activate care

Then install all requirements using:

pip install -r requirements.txt

We have documented detailed steps to help prepare for downloading model checkpoints.

CARE as a CT Reconstruction Enhancement Baseline

Pretrained Autoencoder Checkpoint

huggingface-cli download TianyuLin/CARE --include="autoencoder/*" --local-dir="./STEP1-AutoEncoderModel/klvae/"

Pretrained Diffusion Model Checkpoint

huggingface-cli download TianyuLin/CARE --include="diffusion/*" --local-dir="./STEP2-DiffusionModel/"

Pretrained CARE Model Checkpoints

huggingface-cli download TianyuLin/CARE --include="CARE/*" --local-dir="./STEP3-CAREModel/"

Note

The following script is designed for the nine reconstruction methods mentioned in the paper: three traditional reconstruction methods (FDK, SART, ASD-POCS), five NeRF-based reconstruction methods (InTomo, NeRF, TensoRF, NAF, SAX-NeRF) using the SAX-NeRF Repo, and a Gaussian-Spaltting-based method R2-GS based on its own R2-GS Repo. Feel free to edit to fit your need.

Firstly, Based on the CT reconstruction results from SAX-NeRF Repo and GitHub Repo, please use the provided script to format the dataset:

cd ./ReconstructionPipeline/  # working directory
python -W ignore step1_softlink_BDMAP_O.py   # place the ground truth CT and segmentation
python -W ignore step2_extractAndpixelMetric.py # calculate pixel-wise metrics (SSIM and PSNR)

The resulting dataset format is:

└── BDMAP_O/                      # ground truth folder
    └── BDMAP_O0000001
        └── ct.nii.gz   # the ground truth CT scan of this case
└── BDMAP_O_methodName_numViews/  # reconstruction results folder
    └── BDMAP_O0000001
        └── ct.nii.gz   # the reconstructed CT from `methodName` method with `numViews` X-rays

Run the inference of CARE model via:

cd ./STEP3-CAREModel
bash inference.sh nerf_50 
# e.g. using nerf as CT reconstruction baseline

0. Train Anatomy Segmentator

The anatomy segmentator is built upon nnU-Net, please prepare the dataset first based on nnU-Net's dataset format. Our anatomy segmentator requires 25 anatomical structures. Then, train the model (datasetNum is your actual nnU-Net dataset number):

export nnUNet_raw="/path/to/Dataset_raw"
export nnUNet_preprocessed="/path/to/Dataset_preprocessed"
export nnUNet_results="/path/to/nnUNet_results"
DATASET_NUM=# your actual 3-digit nnU-Net dataset number
nnUNetv2_plan_and_preprocess -d $DATASET_NUM -pl nnUNetPlannerResEncL --verify_dataset_integrity
nnUNetv2_train $DATASET_NUM 3d_fullres all -p nnUNetResEncUNetLPlans # 3d version for anatomy-aware metrics
nnUNetv2_train $DATASET_NUM 2d all -p nnUNetResEncUNetLPlans         # 2d version for CARE model training

1. Train Autoencoder Model

First time Diffusers user should run accelerate config for training device configuration. Assume that the original dataset is in BDMAP format (e.g., AbdomenAtlas 1.0), please use ./ReconstructionPipeline/niigz2h5.py to convert the dataset into *.h5 format.

Next, go to the working directory:

cd ./STEP1-AutoEncoderModel/klvae

Then, modify TRAIN_DATA_DIR and validation_images in train.sh and run the trianing:

bash train.sh

2. Train Diffusion Model

The dataset remained unchanged comparing to the Autoencoder section. Entering the working directory:

cd ./STEP2-DiffusionModel

Then, change the FT_VAE_NAME to select a pretrained autoencoder checkpoint, and run the training:

bash train.sh

3. Train CARE Model

To train the CARE model, first prepare the dataset by:

cd ./ReconstructionPipeline/
python -W ignore step1_softlink_BDMAP_O.py   # place the ground truth CT and segmentation
python -W ignore step2_extractAndpixelMetric.py # format the reconstruction & pixel-wise metrics

Next, go to the working directory:

cd ./STEP3-CAREModel

Then, change the FT_VAE_NAME and TRAINED_UNET_NAME to select a pretrained autoencoder/diffusion model's checkpoint, and run the training:

bash train.sh nerf_50  # example

Citation

@article{lin2025pixel,
  title={Are Pixel-Wise Metrics Reliable for Sparse-View Computed Tomography Reconstruction?},
  author={Lin, Tianyu and Li, Xinran and Zhuang, Chuntung and Chen, Qi and Cai, Yuanhao and Ding, Kai and Yuille, Alan L and Zhou, Zongwei},
  journal={arXiv preprint arXiv:2506.02093},
  year={2025}
}

Acknowledgement

This work was supported by the Lustgarten Foundation for Pancreatic Cancer Research and the Patrick J. McGovern Foundation Award. We would like to thank the Johns Hopkins Research IT team in IT@JH for their support and infrastructure resources where some of these analyses were conducted; especially DISCOVERY HPC DISCOVERY HPC. We thank Hamed Hooshangnejad, Heng Li, Wenxuan Li, and Guofeng Zhang for their helpful suggestions throughout the project.