Deploy generative image priors for image reconstruction using BART

Figure 1. Overview of the proposed method.

Abstract As the image priors are almost always trained in an offline setting using Python, this work aims to deploy the trained model with an MR image reconstruction toolbox, BART, which is a versatile tool for image reconstruction. As shown in Figure 1, there are two steps to realize this: (a) export the constructed computation graph with TensorFlow; (b) use the graph as regularization in BART.

The deployment of models trained with spreco into BART requires minimum environmental prerequisites, providing a practical and user-friendly solution for medical image reconstruction tasks.

Figure 2. The structure of spreco.

Method We developed a Python library for training generative models, spreco, based on TensorFlow, which has the following features:

Distributed training
Interruptible training
Efficient dataloader for medical images
Customizable with a configuration file

And we trained a log-likelihood prior, \(\log p({x};\mathtt{NET}(\hat{\Theta}, {x}))\)\(^1\)\(^{,\ 2}\), which is used to impose learned prior knowledge of images in the SENSE model. The reconstruction is commonly formulated as the following minimization problem \begin{equation} \hat{x}=\underset{x}{\arg\min}\ |\mathcal{A}{x}-{y}|_2^2 + \lambda \log p({x};\mathtt{NET}(\hat{\Theta}, {x})),\nonumber \label{eq:1} \end{equation} where the first term ensures data consistency between the acquired k-space data \({y}\) and the desired image \({x}\), \(\mathcal{A}\) is the forward operator. This problem is solved with FISTA algorithm that has been implemented in BART. To use the learned log-likelihood prior with BART, we have implemented a wrapper using TensorFlow C API for the initialization, restoration and inference of the exported trained model. Click here and give a quick tryout!

Figure 3 displays the evolution of image during the process of reconstruction.

Figure 3. The left sub-figure shows the evolution of the probability density function (PDF) (dashed curves) and the empirically learned PDF (solid curves) at five selected pixels over iterations. The middle sub-figure shows the evolution of magnitude image. The right sub-figure shows the convergence of metrics.

Conclusion We showcased the seamless integration of a TensorFlow trained model into the existing MRI reconstruction workflows of the BART toolbox. This integration allows us to leverage the image reconstruction capabilities offered by BART while harnessing the advantages of rapid prototyping of advanced deep learning models using TensorFlow.

Purpose To develop a deep-learning-based image reconstruction framework for reproducible research in MRI. Methods The BART toolbox offers a rich set of implementations of calibration and reconstruction algorithms for parallel imaging and compressed sensing. In this work, BART was extended by a nonlinear operator framework that provides automatic differentiation to allow computation of gradients. Existing MRI-specific operators of BART, such as the nonuniform fast Fourier transform, are directly integrated into this framework and are complemented by common building blocks used in neural networks. To evaluate the use of the framework for advanced deep-learning-based reconstruction, two state-of-the-art unrolled reconstruction networks, namely the Variational Network and MoDL, were implemented. Results State-of-the-art deep image-reconstruction networks can be constructed and trained using BART’s gradient-based optimization algorithms. The BART implementation achieves a similar performance in terms of training time and reconstruction quality compared to the original implementations based on TensorFlow. Conclusion By integrating nonlinear operators and neural networks into BART, we provide a general framework for deep-learning-based reconstruction in MRI.

Deploy generative image priors for image reconstruction using BART

References

2023

2022

2021