Official code for Score-Based Generative Models (SGM) for PET Image Reconstruction (MELBA, accepted) by Imraj RD Singh, Alexander Denker, Riccardo Barbano, Željko Kereta, Bangti Jin, Kris Thielemans, Peter Maass, Simon Arridge.
I. Singh, A. Denker and R. Barbano have equal contribution.
In this work we address PET-specific challenges such as; non-negativity of measurements/images, varying dynamic range of underlying radio-tracer distributions, and low-count Poisson noise on measurements requiring a Poisson Log-Likelihood (PLL). Further, we develop methods for 3D reconstruction, propose a guided variant with a Magnetic Resonance (MR) image, and accelerate the method using subsets.
Our modifications can be summarised with the following diagram:
Where the sections pertain to those in the paper. The most appropriate reconstruction proposed, PET-variant of Decomposed Diffusion Sampling (PET-DDS; where DDS is proposed for MRI and CT here), was extended to 3D and the reconstruction steps are illustrated below:
The work presented develops and adopts code from various repositories, where specific contributions are indicated at the top of sources. The most important repositories include:
- SGM sampling methods for inverse problems
- pyParallelProj for 2D experiments data generation
- SIRF-exercises for 3D experiments data generation
- Normalised supervised PET baselines
- DIVal for supervised deep learning architectures and training scripts
- Guided diffusion repository for the diffusion model architecture
- Deep image prior comparison
We thank the authors of the aforementioned repositories for their open-source development and contributions.
The results of this work are in-silico simulations of the BrainWeb dataset, and all datasets are freely available for download/generation. For 2D work, and training the score-model, we use the dataset available here, which can be downloaded through pyParalellProj. For 3D work we use the dataset available here here.
Files for the generation of 2D data can be found in src/brainweb_2d/. For 3D data generation we provide a juypter notebook src/sirf/brainweb_3D.ipynb.
Training of the score-model requires running script main_score_based_models_train.py. All experiments with reconstruction techniques can be found in coordinators/, and all results can be processed with files in results/.
For reproducibility we provide a devcontainer utilising docker to containerise the development environment required for this work. The files are located in .devcontainer/, these files use scripts to setup up conda environments where the environment is defined with files in scripts/, we provided full list of static dependencies in req.txt. Please note that this project requires SIRF for 3D work.
Arxiv bibtex:
@article{melba:2024:001:singh,
title = "Score-Based Generative Models for PET Image Reconstruction",
author = "Singh, Imraj RD and Denker, Alexander and Barbano, Riccardo and Kereta, Željko and Jin, Bangti and Thielemans, Kris and Maass, Peter and Arridge, Simon",
journal = "Machine Learning for Biomedical Imaging",
volume = "2",
issue = "Special Issue for Generative Models",
year = "2024",
pages = "547--585",
issn = "2766-905X",
doi = "https://doi.org/10.59275/j.melba.2024-5d51",
url = "https://melba-journal.org/2024:001"
}