/DiffPose

[CVPR 2024] Intraoperative 2D/3D registration via differentiable X-ray rendering

Primary LanguagePythonMIT LicenseMIT

DiffPose

Intraoperative 2D/3D registration via differentiable X-ray rendering

CI Paper shield License: MIT Docs Code style: black

Install

To install DiffPose and the requirements in environment.yml, run:

pip install diffpose

The differentiable X-ray renderer that powers the backend of DiffPose is available at DiffDRR.

Datasets

We evaluate DiffPose networks on the following open-source datasets:

Dataset Anatomy # of Subjects # of 2D Images CTs X-rays Fiducials
DeepFluoro Pelvis 6 366
Ljubljana Cerebrovasculature 10 20
  • DeepFluoro (Grupp et al., 2020) provides paired X-ray fluoroscopy images and CT volume of the pelvis. The data were collected from six cadaveric subjects at John Hopkins University. Ground truth camera poses were estimated with an offline registration process. A visualization of one X-ray / CT pair in the DeepFluoro dataset is available here.
mkdir -p data/
wget --no-check-certificate -O data/ipcai_2020_full_res_data.zip "http://archive.data.jhu.edu/api/access/datafile/:persistentId/?persistentId=doi:10.7281/T1/IFSXNV/EAN9GH"
unzip -o data/ipcai_2020_full_res_data.zip -d data
rm data/ipcai_2020_full_res_data.zip
  • Ljubljana (Mitrovic et al., 2013) provides paired 2D/3D digital subtraction angiography (DSA) images. The data were collected from 10 patients undergoing endovascular image-guided interventions at the University of Ljubljana. Ground truth camera poses were estimated by registering surface fiducial markers.
mkdir -p data/
wget --no-check-certificate -O data/ljubljana.zip "https://drive.google.com/uc?export=download&confirm=yes&id=1x585pGLI8QGk21qZ2oGwwQ9LMJ09Tqrx"
unzip -o data/ljubljana.zip -d data
rm data/ljubljana.zip

Experiments

To run the experiments in DiffPose, run the following scripts (ensure you’ve downloaded the data first):

# DeepFluoro dataset
cd experiments/deepfluoro
srun python train.py     # Pretrain pose regression CNN on synthetic X-rays
srun python register.py  # Run test-time optimization with the best network per subject
# Ljubljana dataset
cd experiments/ljubljana
srun python train.py
srun python register.py

The training and test-time optimization scripts use SLURM to run on all subjects in parallel:

  • experiments/deepfluoro/train.py is configured to run across six A6000 GPUs
  • experiments/deepfluoro/register.py is configured to run across six 2080 Ti GPUs
  • experiments/ljubljana/train.py is configured to run across twenty 2080 Ti GPUs
  • experiments/ljubljana/register.py is configured to run on twenty 2080 Ti GPUs

The GPU configurations can be changed at the end of each script using submitit.

Development

DiffPose package, docs, and CI are all built using nbdev. To get set up withnbdev, install the following

conda install jupyterlab nbdev -c fastai -c conda-forge 
nbdev_install_quarto      # To build docs
nbdev_install_hooks       # Make notebooks git-friendly
pip install -e  ".[dev]"  # Install the development verison of DiffPose

Running nbdev_help will give you the full list of options. The most important ones are

nbdev_preview  # Render docs locally and inspect in browser
nbdev_clean    # NECESSARY BEFORE PUSHING
nbdev_test     # tests notebooks
nbdev_export   # builds package and builds docs
nbdev_readme   # Render the readme

For more details, follow this in-depth tutorial.

Citing DiffPose

If you find DiffPose or DiffDRR useful in your work, please cite the appropriate papers:

@article{gopalakrishnan2023intraoperative,
  title={Intraoperative {2D/3D} Image Registration via Differentiable X-ray Rendering},
  author={Gopalakrishnan, Vivek and Dey, Neel and Golland, Polina},
  journal={arXiv preprint arXiv:2312.06358},
  year={2023}
}

@inproceedings{gopalakrishnan2022fast,
  title={Fast Auto-Differentiable Digitally Reconstructed Radiographs for Solving Inverse Problems in Intraoperative Imaging},
  author={Gopalakrishnan, Vivek and Golland, Polina},
  booktitle={Workshop on Clinical Image-Based Procedures},
  pages={1--11},
  year={2022},
  organization={Springer}
}