Differentiable Score-Based Likelihoods: Learning CT Motion Compensation From Clean Images

This repository contains the source code for our paper Differentiable Score-Based Likelihoods: Learning CT Motion Compensation From Clean Images which has been presented at MICCAI 2024. Trained weights and example data are provided to run the proposed CT motion algorithm using differentiable likelihood computation of diffusion models, as well as code and weights for the reference methods. This repository also contains the code to create the full data set based on the publicly available CQ500 data set which can be used to train the models.

Usage with pretrained weights

This section describes how to run the motion compensation on an example head CT slice with pretrained weights.

1. Clone this repository including submodules

   git clone https://github.com/mareikethies/moco_diff_likelihood.git --recurse_submodules
   

2. Download network weights and example data from figshare using the following links: Example data, diffusion model weights, autofocus model weights. Place both model weights unchanged in the checkpoints subfolder. Extract the example data and place it in the data subfolder. The folder structure should look like

   ./data/CQ500CT268/image_60/
   

3. To apply a random motion pattern to the example data and compensate it...

   ...using the proposed diffusion model likelihood target function (requires ~ 10 GB GPU memory), run

   python motion_compensation.py
   

   ...using the autofocus target function, run

   python motion_compensation.py target_function.choice='autofocus' optimizer.options.stepsize_translation=50 optimizer.options.stepsize_rotation=50
   

   ...using the supervised MSE optimization with ground truth, run

   python motion_compensation.py target_function.choice='mse_supervised' optimizer.options.stepsize_translation=5000 optimizer.options.stepsize_rotation=5000
   

   Results will be stored in the out subfolder.

Recreate full data set and train models

This section describes how to process the full CQ500 data set to reproduce the results from the paper and how to use it for training the networks.

1. Check the license and download the CQ500 head CT data set from here.

2. Filter the data set for reconstructions with thin slice thickness using clean_cq500.py.

3. Create a fan-beam data set using generate_data.py.

4. To train the score-based diffusion model, run

   python score_sde_pytorch/main.py --config motion_ncsnp_continuous_small.py --mode train --workdir your_output_path --datadir your_path/HeadSimulatedProjectionDataCQ500FanBeam/motion_free/
   

5. To train the autofocus model, run

   python autofocus_train.py --data_dir your_path/HeadSimulatedProjectionDataCQ500FanBeam/motion_free --run_name your_name --accelerator cuda --devices 1 --max_epochs 2000 --batch_size 32 --learning_rate 1e-3 --num_workers 8
   

The full data set created in steps 1 - 3 instead of the example data and the image_lookup_paper_results.json instead of image_lookup.json, will recreate the results in the published paper.

Dependencies and licenses

This code depends on (1) our previously published work on gradient propagation to CT geometry parameters (see paper, see github) and (2) the original work on score-based diffusion models by Song et al. (see paper in particular section D.2 of the appendix, see github). The source code for both works is published under Apache 2.0 license.

We further provide sources for the PyroNN package as used for development of this project. The original, but outdated repository is here. For this project, please install pyronn from the provided sources in the pyronn folder by running (from within that folder)

python -m pip install -e pyronn-layers-torch
python -m pip install -e pyronn

This requires a version of pytorch < 2.0.

Citation

If you use this code for your research, please cite our paper:

@InProceedings{10.1007/978-3-031-72104-5_25,
    author="Thies, Mareike and Maul, Noah and Mei, Siyuan and Pfaff, Laura and Vysotskaya, Nastassia and Gu, Mingxuan and Utz, Jonas and Possart, Dennis and Folle, Lukas and Wagner, Fabian and Maier, Andreas",
    editor="Linguraru, Marius George and Dou, Qi and Feragen, Aasa and Giannarou, Stamatia and Glocker, Ben and Lekadir, Karim and Schnabel, Julia A.",
    title="Differentiable Score-Based Likelihoods: Learning CT Motion Compensation from Clean Images",
    booktitle="Medical Image Computing and Computer Assisted Intervention -- MICCAI 2024",
    year="2024",
    publisher="Springer Nature Switzerland",
    address="Cham",
    pages="253--263",
}

If you have any questions about this repository or the paper, feel free to reach out ([email protected]).