AXON: Conditional Diffusion for 3D CT Volume Reconstruction from 2D X-rays

🎉 Our paper was accepted at BMVC 2026!

Official implementation of

Conditional Diffusion for 3D CT Volume Reconstruction from 2D X-rays — BMVC 2026

AXON (Advanced X-ray to CT-volume Network) reconstructs a high-resolution 3D CT volume from one or two real 2D chest X-rays with a three-stage cascade of generative models:

• AXON-CoarseDiff (CoarseDiff/) — a 3D Brownian Bridge diffusion model (BBDM) that lifts the X-ray projections into a coarse CT volume.
• AXON-FineDiff (FineDiff/) — a rectified-flow 3D diffusion UNet with a ControlNet conditioned on the CoarseDiff output, refining anatomy and intensities.
• AXON-SR (SR/) — a 3D RRDB (ESRGAN-style) network that super-resolves the refined volume from 128³ to 256³.
• Generation pipeline (Pipeline/) — end-to-end inference chaining all three stages.

AXON supports conditioning on one view (frontal X-ray only) or two views (frontal + lateral). This is a single switch — data.views in configs/common.yaml — that consistently controls the datalist, the dataloaders, the CoarseDiff conditioning branch (with two views, both projections are encoded separately and merged by a learned 3D fusion block) and the generation pipeline. There is only one implementation of each stage; nothing is duplicated per view setting.

Repository layout

AXON/
├── configs/
│   ├── common.yaml               # shared parameters for ALL stages (see below)
│   └── datalists/                # LIDC-IDRI datalists (1-view and 2-view)
├── common/                       # loader for the shared config
├── CoarseDiff/                   # stage 1: training + inference (1 or 2 views)
├── FineDiff/                     # stage 2: diffusion UNet + ControlNet
├── SR/                           # stage 3: super resolution
├── Pipeline/                     # full inference pipeline (reuses the
│                                 # CoarseDiff/SR implementations, 1 or 2 views)
├── Figures/                      # figures used in this README
├── environment.yml               # single conda environment for everything
└── requirements.txt

Installation

One environment covers all stages:

conda env create -f environment.yml
conda activate axon

or, with an existing Python ≥ 3.10 environment:

# install a PyTorch build matching your CUDA version first, e.g.
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu124
pip install -r requirements.txt

Shared configuration

Everything the three stages have in common — the number of views, dataset root, datalist paths, validation fold, volume/X-ray resolutions, and the HU clipping range — lives in configs/common.yaml and is injected into every training and inference entry point automatically (each script also accepts --common_config <path> to point at an alternative file). The per-stage configs (CoarseDiff/configs/coarsediff.yaml, FineDiff/configs/*.json, SR/configs/sr.yaml, Pipeline/configs/*) only contain stage-specific hyperparameters.

• data.views — 2 conditions on a frontal + lateral X-ray pair, 1 on the frontal X-ray only. All stages and the pipeline follow this switch; models must be trained and evaluated with the same setting.
• Defaults: volumes are processed at 128³ (SR output 256³), X-rays at 1024², CT intensities are clipped to [-100, 900] HU and normalized to [-1, 1]. Fold 2 of the datalists is always the test fold; data.fold selects the validation fold.

Data preparation

The datalists in configs/datalists/ describe the preprocessed LIDC-IDRI dataset; data.views selects between the two-view and one-view list, which reference the same files on disk and use identical folds. Set data.base_dir in configs/common.yaml to the folder that contains:

File	Content
<id>-CT.h5	CT volume as HDF5 with datasets volume (H×W×D), affine (4×4) and spacing
<id>-X-RAY-FRONT.tiff	frontal X-ray or DRR projection (float TIFF)
<id>-X-RAY-SIDE.tiff	lateral X-ray or DRR projection (float TIFF; 2-view setting only)
bbdm_2_views_output/<id>_syn_CT.nii.gz	2-view CoarseDiff outputs (only needed to train FineDiff, see below)
bbdm_one_view_output/<id>_syn_CT.nii.gz	1-view CoarseDiff outputs (only needed to train FineDiff with views: 1)

Datalist structure

A datalist is a JSON file with a single training list; every entry describes one case. Users are expected to bring their own data and build the corresponding datalists — the files shipped in configs/datalists/ show the expected format for LIDC-IDRI. A two-view entry looks like this:

{
  "id": "LIDC-IDRI-0001",
  "image_original": "LIDC-IDRI-0001-CT.h5",
  "image": "LIDC-IDRI-0001-CT.h5",
  "label_front": "LIDC-IDRI-0001-X-RAY-FRONT.tiff",
  "label_side": "LIDC-IDRI-0001-X-RAY-SIDE.tiff",
  "bbm_output": "bbdm_2_views_output/LIDC-IDRI-0001_syn_CT.nii.gz",
  "fold": 1,
  "spacing": [1.0, 1.0, 1.0]
}

• All paths are relative to data.base_dir.
• The one-view datalist uses a single label key (frontal view) instead of label_front/label_side.
• bbm_output points to the CoarseDiff output of the case and is only required for FineDiff training.
• fold controls the split: fold 2 is always the test fold, the fold selected by data.fold in configs/common.yaml is the validation fold, and all remaining folds are used for training.
• spacing must be present but its value is not used: every stage resamples volumes to 1.0×1.0×1.0 mm³ and then substitutes unit spacing, so [1.0, 1.0, 1.0] is the honest placeholder.
• MAISI's body-region conditioning is disabled in AXON (include_body_region: false in FineDiff/configs/config_network.json), so no top_region_index / bottom_region_index entries are needed. The FineDiff trainers raise an explicit error if that flag is turned on.

Training entries can freely mix cases with real X-rays and cases with synthetic DRR projections (e.g. generated with DeepDRR); validation and testing should use real X-rays. CT volumes are resampled to 1.0×1.0×1.0 mm³ before rescaling to 128³, and real X-rays are registered to their DRR reference (see the preprocessing section of the paper).

Training

Train the stages in order — FineDiff is conditioned on CoarseDiff outputs. Set data.views in configs/common.yaml (1 or 2) before training; all stages of one model family must be trained with the same view setting.

Stage 1: CoarseDiff

cd CoarseDiff
bash train.sh          # python3 main.py --config configs/coarsediff.yaml --train ...

Checkpoints and samples go to CoarseDiff/results/LIDC_IDRI/coarsediff/. To evaluate on the test fold:

bash test.sh

Generating the FineDiff training data. FineDiff's ControlNet is conditioned on CoarseDiff outputs for all cases, whereas test.sh only evaluates the validation fold and writes to CoarseDiff/results/<dataset>/<model>/sample_to_eval/. Use the dedicated script instead:

bash generate_finediff_data.sh

It runs CoarseDiff over the training, validation and test split in one pass and writes <id>_syn_CT.nii.gz straight to <data.base_dir>/bbdm_2_views_output/ (2-view setting) resp. <data.base_dir>/bbdm_one_view_output/ (1-view setting) — the locations referenced by the bbm_output entries of the datalists. Cases whose output already exists are skipped, so an interrupted run can be restarted; use bash generate_finediff_data.sh --overwrite to regenerate them, and --splits train,val to restrict the run.

Stage 2: FineDiff

FineDiff consists of two trainings. First the diffusion UNet (on CT volumes only), then the ControlNet (conditioned on the CoarseDiff outputs) with the UNet frozen:

cd FineDiff
bash train_diff_unet.sh     # stage 2a: diffusion UNet
bash train_controlnet.sh    # stage 2b: ControlNet (needs 2a + CoarseDiff outputs)

Checkpoints are written to FineDiff/models/ (*_best.pt tracks the best validation PSNR). Paths such as the frozen UNet checkpoint are set in FineDiff/configs/environment_*.json. Evaluate the ControlNet on the test fold with:

bash test_controlnet.sh

Stage 3: SR

The SR network is trained on LR/HR pairs created from the ground-truth CTs (128³ inputs obtained by downsampling, 256³ targets):

cd SR
bash train.sh          # python3 train.py --config configs/sr.yaml

Checkpoints go to SR/outputs/sr/checkpoints/. For test-fold evaluation, set checkpoint_path in SR/configs/sr.yaml and run bash test.sh.

Inference: full generation pipeline

The pipeline runs all three stages on the test fold, reports MAE / MSE / PSNR / SSIM after every stage, and can save the generated volumes as NIfTI. It contains no model code of its own — the CoarseDiff model/runner and the SR network definition are imported from their stage directories, so training and inference always share one implementation.

Make sure data.views in configs/common.yaml matches the setting the models were trained with, then run:

cd Pipeline
python3 main.py \
    --config_bbm configs/pipeline.yaml \
    --gpu_id 0 \
    --resume_model_bbm <coarsediff checkpoint (.pth)> \
    --finediff_network_config configs/config_network.json \
    --finediff_train_config configs/config_train.json \
    --finediff_unet_ckpt <finediff diffusion-UNet checkpoint (.pt)> \
    --finediff_controlnet_ckpt <finediff ControlNet checkpoint (.pt)> \
    --finediff_sampling_steps 100 \
    --sr_config configs/sr.yaml \
    --sr_ckpt <sr checkpoint (.pt)> \
    --log_metrics ./results/pipeline_metrics.txt \
    --save_nifti \
    --result_path ./results

(run_pipeline.sh contains the same call as an editable template.)

Per-sample metrics are written to results/per_sample_metrics.csv, aggregated metrics to results/pipeline_metrics.txt, and generated volumes to results/<id>_{coarsediff,finediff,sr}.nii.gz.

Example results

Reconstructions from a real LIDC-IDRI X-ray compared to prior methods and the ground truth (axial, sagittal, and coronal mid-slices with 3D renderings):

Acknowledgements

This repository builds directly on the excellent work and code of others — many thanks to their authors:

• Cor2Vox (Bongratz et al.) — AXON-CoarseDiff is based on the Cor2Vox code base, which itself builds on BBDM (Li et al.), the Brownian Bridge Diffusion Model.
• MAISI (NVIDIA / Project MONAI) — AXON-FineDiff adapts the MAISI diffusion UNet and ControlNet implementation and training scripts.
• ESRGAN (Wang et al.) — AXON-SR uses a 3D adaptation of the ESRGAN RRDB architecture.

We also thank the creators of the LIDC-IDRI dataset and DeepDRR.

License

This repository contains code only — no pretrained weights are distributed, and none are required: every stage is trained from scratch as described in the Training section. The subdirectories retain the licenses of the projects the code is derived from (see CoarseDiff/LICENSE, SR/LICENSE, and the Apache-2.0 headers in the MONAI-derived files under FineDiff/).

Citation

If you use this code, please cite the AXON paper:

@inproceedings{rath2026axon,
  title     = {Conditional Diffusion for 3D CT Volume Reconstruction from 2D X-rays},
  author    = {Rath, Martin and Ghahremani, Morteza and Li, Yitong and
               Taghipour, Ashkan and Makowski, Marcus and Wachinger, Christian},
  booktitle = {British Machine Vision Conference (BMVC)},
  year      = {2026}
}