Conversion Between CT and MRI Images Using Diffusion and Score-Matching Models
Qing Lyu, Ge Wang
I Introduction
Magnetic resonance imaging (MRI) and computed tomography (CT) are two most widely used medical imaging modalities. MRI shows soft tissues such as vessels and organs in rich contrast while CT is the method of choice for imaging hard tissues like bones as well as interfaces among air, bone and soft tissues. Due to their complementary characteristics, multi-modality imaging with MRI and CT is often used in clinical practice. For example, radiotherapy requires both CT and MRI, because CT provides an electron density distribution indispensable for treatment planning while MRI outlines tumors and soft tissues . However, simultaneous CT-MRI is still a research area, and CT and MRI scans are currently performed in separation, which is not only expensive but also brings about nonrigid misalignment between MRI and CT images . Developing a simultaneous CT-MRI device may be a solution of this problem, and we have conducted studies to propose top-level designs of such a device .
Medical image synthesis is a viable approach to solve the aforementioned problem. This approach models a mapping from a given source image to an unknown target image. Conventional image synthesis methods focus on exploiting diverse models, such as dictionary learning and random forest, to extract features predefined by experts . However, these methods are limited to handcrafted feature representations. Recently, deep learning has shown huge potentials and great successes in medical image analysis tasks, such as denoising , super-resolution , and artifact reduction . Compared with the traditional methods, deep neural networks learn features in a data-driven fashion and produce superior feature representations. In particular, several deep learning-based cross-modality medical image synthesis studies were reported in recent years, most of which are based on convolutional neural networks (CNNs) and generative adversarial networks (GANs) .
Diffusion and score-matching models represents an emerging generative approach that has attracted a major attention in the medical imaging field. These models can generate high-fidelity realistic natural images . Compared with other types of generative models like GANs and variational autoencoders that are difficult to train and interpret, and do not always produce satisfactory image quality, diffusion and score-matching models are analytically principled, and easy to train, and offer state-of-the-art image quality. Impressively, an increasing number of studies show that diffusion and score-matching models beat GANs and variational auto-encoders in multiple image generation tasks .
In this paper, we propose to use diffusion and score-matching models for image conversion between CT and MRI, with an emphasis on mapping from MRI to CT images. Two models are based on in our study, the denoising diffusion probabilistic model (DDPM) and the model solving the stochastic differential equation (SDE) . These models are compared with CNN and GAN models of consistent network architectures. With the diffusion models, we further quantify their uncertainties from Monte-Carlo sampling results and obtain superior results by averaging these random samples.
II Related works
Inspired by the success of deep learning in the computational vision domain, researchers applied deep learning in medical imaging tasks. For medical image synthesis, CNN and GAN models were proposed. For image conversion between MRI and CT modalities, Nie et al. proposed a 3D fully convolutional network to synthesize pelvic CT images from the corresponding MRI images. Then, they proposed a cascaded GAN model with multiple sequential generators and discriminators . Bahrami et al. designed a CNN model based on an encoder-decoder backbone and reported that the proposed model exhibited a fast convergence rate with a low number of training subjects. Han et al. and Leynes et al. both synthesized CT images using UNet. GAN models with ResNet were created by Emami et al. and Tao et al. . Boni et al. built a conditional GAN to generate synthetic CT images based on multi-center pelvic datasets. Furthermore, Chartsias et al. , Hiasa et al. , Zhang et al. , and Cai et al. adopted CycleGAN to generate synthetic MRI and CT images. Li et al. compared the performance of UNet, cycleGAN, and pix2pix models on converting brain CT scans into MRI images, and found that UNet had the best performance among all the three models.
For image synthesis between CT and positron emission tomography (PET) modalities, Ben-Cohen et al. proposed a conditional GAN with a fully convolutional network for liver PET image synthesis. Armanious et al. built a model called MedGAN with cascaded encoder-decoders to minimize a complicated objective function. Bi et al. designed a multi-channel GAN model with a ability to represent semantic information. Ben-Cohen et al. used a GAN model with a fully connected network to generate synthetic PET images and improve lesion detection.
As far as image synthesis between MRI and PET modalities is concerned, Wei et al. proposed a sketcher-refiner scheme with two cascaded GANs. The first GAN generates coarse synthetic images. The second GAN refines the results. Choi et al. built a GAN model with UNet as the generator for MRI image synthesis. Zhang et al. proposed a model called BPGAN to synthesize brain PET images. Hu et al. designed a 3D end-to-end synthesis model called the bidirectional mapping generative adversarial network (BMGAN), in which the image context and the latent vector were jointly optimized for brain MRI-to-PET image synthesis.
II-B Diffusion and score-matching model
Diffusion and score-matching models are emerging as most promising deep generative models, with impressive generative capabilities for many tasks such as image generation, super-resolution, and image inpainting . Such a model usually consists of two stages: a forward stage to gradually add noise, and a reverse stage to denoise and recover an original sample step-by-step. Currently, representative frameworks in this category of image generation methods include denoising diffusion probabilistic models (DDPM) , noise conditioned score networks (NCSN) , and stochastic differential equations (SDE) .
DDPM has its diffusion stage including multiple small steps. In each step, a data sample such as an image is slightly corrupted by Gaussian noise. Let represents an original image and denotes the original distribution of , we have . A sequence of gradually corrupted images after each diffusion step can be computed as the following Markovian process:
where is the total number of noising steps, is a hyper-parameter controlling the variance of incremental Gaussian noise, and represents a Gaussian distribution of mean and covariance . With the parametrization and , we have
When , becomes an isotropic Gaussian distribution.
In its reverse stage, DDPM performs a denoising task to recover an original image. According to , each step is also a Gaussian distribution if is small. Then, we can train a neural network to approximate each reserve diffusion step and estimate the mean and the covariance :
where is the density function of . According to , the reverse step is tractable conditioned on and :
Given that , where . Equation (7) can be rewritten as
The objective for training the noise estimation network (added noise in ) is to optimize the variational lower bound (VLB):
where denotes the Kullback-Leibler divergence between two probability distributions. Note that is constant and can be ignored because has no learnable parameters and is a Gaussian noise.
In , is computed from , and the loss term in (12) can be reparameterized and simplified as
The final simplified objective function is
where is a constant independent of the vector of parameters .
II-B2 Noise conditioned score network
Langevin dynamics produces samples from a probability density function only using the score function . Given a fixed step size and an initial value with being a prior distribution, the sampling process using the Langevin method can be expressed as
where . The distribution of equals when and . A neural network is trained to estimate the score so that . Ideally, the network can be trained via score matching based on the following objective function
However, equation (17) is hard to be optimized as the score is not easy to obtain. To overcome this difficulty, Song et al. proposed to perturb the original data distribution by Gaussian noises at different scales: such that and . Then, a NCSN is trained for score estimation: . Then, we have
Combining (17) and (18) on all , we have
where is a weighting factor.
II-B3 Stochastic differential equation
Similar to DDPM and NCSN, the stochastic differential equation (SDE) framwork gradually transforms the original data distribution into a Gaussian distribution in the forward stage. Unlike the other two methods that split the diffusion process into many discrete steps, the SDE method handles a continuous process. Thus, the SDE method can be seen as a generalization of DDPM and NCSN methods. Let us use for the probability density function of , and for the transition kernel from to , where . Typically, is an unstructured prior distribution without information from . To calculate the diffusion process of the SDE method, it is necessary to solve the following SDE
where , is the Brownian motion, and are drift and diffusion coefficients respectively. Similarly, we have the associated reverse-time SDE:
where is the score function of the data distribution , represents the Brownian motion when time is reversed, and . A neural network is trained to estimate the score so that . The objective function is the continuous version of (19) and can be expressed as
where is a positive weighting function, , , and . In (22), is used to replace the original score, as proposed in .
III Methodology
The Gold Atlas male pelvis dataset was utilized in this study, which consists of co-registrated T2w MRI and CT image pairs from 19 patients. Data were collected in three different departments. CT images were obtained on a Siemens Somantom Definition AS+ scanner, a Toshiba Aquilion scanner, and a Siemens Emotion 6 scanner with the pixel size range between and . T2w MR images were scanned on a GE Discovery 750w scanner with the FRFSE sequence, a Siemens scanner with the TSE sequence, and a GE signa PET/MR scanner with a FRFSE sequence, with the pixel size range between and . We randomly selected 17 patients with 1,416 image pairs for training, and the other two patients with 135 image pairs for testing. All image matrices in this study are . Before training, all images were pre-processed for pixel intensity unification.
III-B Conditional DDPM
To implement image synthesis between CT and MRI, here we extend the conditional DDPM proposed by Saharia et al. from working within the same imaging mode (photographs) to mapping across different imaging modes (CT and MRI) so that we can build diffusion and score-matching models conditioned on T2w images. Given the co-registered CT and T2w MRI pairs , where is the number of image pairs in the dataset, our objective function of (14) is as follows:
The sampling process is a reverse Markovian process starting from a Gaussian noise , the reverse process of (6) and (9) can be modified as
The training and sampling procedures of modified DDPM are listed in Table I. There were 1,000 diffusion steps and 1,000 sampling steps for the modified DDPM. UNet was adopted for the reverse diffusion process to denoise.
III-C Conditional SDE
To generate required sampling results, the reverse-time SDE should be solved under the guidance of a condition of interest. There are various approaches for enforcing a condition gently or strongly, such as classifier-free and classier-guidance methods. The classifier-free approach, as mentioned by Song et al. , adds a condition in the diffusion model training process and performs the network training in a supervised fashion. Different from the classifer-free approach, the classifier-guidance approach is unsupervised. Song et al. trained a network for unconditional score estimation and then incorporated conditional information into the sampling process with a proximal optimization step based on a physical measurement model for medical imaging such as the Radon and Fourier transforms. Dharwal et al. and Liu et al. added a classifier and used its gradients to regulate the reverse diffusion process of a pre-trained unconditional diffusion model.
In this study, we include T2w MRI images as the condition in the training process. In other words, we supervise the forward and backward diffusion processes. Specifically, we adopted the variance exploding (VE) SDE setting described in the paper with and . Given the original expression of and , we have . Hence, equation (20) can be changed to the following form:
where . As is a Gaussian perturbation kernel, the gradient of the perturbation kernel is . Accordingly, the objective function becomes
The reverse-time SDE of (21) can be expressed as
To sample from the time-dependent score-based model , we first draw a sample from the prior distribution , and then solve the reverse-time SDE numerically. When is large, the mean value of the prior distribution is close to 0, and we can approximate the prior distribution to . In this study, three sampling techniques were used, which are Euler-Maruyama (EM), Prediction-Corrector (PC), and probability flow ordinary differential equation (ODE) methods respectively.
In the EM method, to solve the reverse-time SDE of (29), a simple discretization strategy is adopted, replacing with a small increment and with a Gaussian noise . Then, we have
where .
III-C2 Prediction-Correction method
The PC sampling alternates between prediction and correction steps. The predictor can be any numerical solver for the reverse-time SDE with a fixed discretization strategy, such as the EM method of (30). The corrector can be any score-based Markov Chain Monte Carlo method, such as annealed Langevin dynamics. To implement annealed Langevin dynamics, it is necessary to calculate a Langevin step size :
where is a signal-to-noise ratio, and . Once the Langevin step size is determined, we can sample according to Langevin dynamics of (16).
III-C3 Probability flow ODE method
We call the probability flow ODE method as ODE in this paper for simplicity. For any SDE in the form of (20), there exists an associated ODE
which has the same marginal probability density trajectory as that of the SDE. As a result, sampling by solving the reverse-time SDE is equivalent to solving the above ODE in the reverse time direction. Being the same as the EM and PC methods, the ODE sampling process starts from obtaining from . Then, we integrate ODE in the reverse time direction and finally get a sample from . In this case, the ODE equation is written as follows:
We used the Explicit Runge-Kutta method of order 5(4) method to solve (33).
The training and sampling procedures of the three mentioned methods are listed in Table II. We used UNet for score estimation. For all the compared sampling methods, we set the total sampling steps to 1,000. In particular, the PC sampling was set with 500 prediction steps and 500 correction steps.
III-D Other methods for comparison
We also compared the diffusion and score-matching models with CNN and GAN based models. For CNN, UNet was trained to minimize MSE. For GAN, Wasserstein distance and gradient penalty were incorporated (WGAN-GP) with UNet as the generator and MSE as data fidelity measure.
III-E Implementation details
In our experiments, for all the methods including conditional DDPM, conditional SDE, CNN, and WGAN-GP, we consistently used UNet of the same architecture (UNet of CNN and WGAN-GP without time-embedding) and the Adam optimizer with a learning rate of , betas of 0.9 and 0.999, and eps of . The training process continued at least 100 epochs and stopped after the loss did not decrease by 1% relative to the average loss for the twenty previous epochs. The batch size was fixed at 2. All the experiments were conducted using PyTorch 1.11 on a 24GB Nvidia RTX Titan GPU. We will upload all the codes for this study onto Github after the publication of this paper.
For the conditional DDPM and conditional SDE methods, since noises were involved in the sampling processes, the final outputs were subject to random fluctuations. Hence, we further investigated uncertainties of the diffusion and score-matching models using the Monte Carlo (MC) method. Each target CT image was generated ten times. In each case, we recorded all sampling results and obtained a MC result by averaging all ten results. We denote these averaged results using DDPM, ODE, EM, and ODE sampling methods as DDPM-M, ODE-M, EM-M and PC-M respectively. The model uncertainty was revealed by looking at the standard deviation map of the ten sampling results in each configuration.
IV Results
In this study, we utilized the structural similarity index measure (SSIM) and peak signal-to-noise ratio (PSNR) metrics to evaluate image quality.
Intermediate results from DDPM, EM, PC, and ODE methods in the reverse process are compared in Fig. 1. It can be found that all the four methods finally generate a desirable CT image () starting from a Gaussian noise () conditioned on the associated T2w MR image. As the reverse diffusion process goes on, noise is gradually removed to make an image increasingly realistic.
IV-B Model uncertainty estimation
Fig. 2 shows MC results using different sampling methods respectively. It seems that DDPM generates results with the highest SSIM and PSNR scores while the ODE method produces results with the lowest scores. For those results obtained by averaging all ten MC samples conditioned on the same T2w MR image, ODE, EM, and PC methods offer significantly better synthetic CT images with higher SSIM and PSNR scores than the corresponding individual result. In terms of the standard deviation map, the PC and EM methods generate lower standard deviations than the DDPM and ODE methods. Quantitative results in Fig. 3(b) shows model uncertainty scores for the sampling methods, which were computed by averaging over a standard deviation map, demonstrating that EM and PC have lower of model uncertainty scores than DDPM and ODE.
IV-C Comparison with CNN and GAN
We compared diffusion and score-matching models with CNN and WGAN-GP based models on image synthesis between CT and MRI. Qualitative results are shown in Fig. 4. It is found that the CNN method tends to generate over-smoothed results. Although the WGAN-GP method generates results with details, it tends to generate artifacts. In the bottom row of Fig. 4, the left side of the round bone structure produced by WGAN-GP contains severe artifacts. On the other hand, the diffusion and score-matching results have faithful details. Among the diffusion and score-matching results, the ODE method generates less favorable results than the other three sampling methods, and DDPM and PC methods show higher SSIM scores than the other two sampling methods. Looking at the averaged results (DDPM-M, ODE-M, EM-M, and PC-M), it is found that all four kinds of averaged results are better than the corresponding individual sampling results in terms of information fidelity, SSIM and PSNR scores. Quantitative results in Fig. 3(a) also demonstrate higher scores of DDPM-M, ODE-M, EM-M, and PC-M in terms of SSIM and PSNR than DDPM, ODE, EM, and PC respectively. We also investigated inference speed of each sampling method. Fig. 3(c) compares the amount of time to generate a synthetic CT image using each sampling method. The ODE method is the fastest among all the four methods. The sampling time of DDPM and EM are comparable while the PC method is the slowest.
V Discussions and conclusion
In this study, DDPM is time-discrete in both training and sampling processes. In contrast, SDE is time-continuous in the training process and time-discrete in the sampling process. We investigated both DDPM and SDE methods generating synthetic CT images from given T2w MRI. Four different sampling methods (one DDPM based and three SDE based) were compared. According to our results, all the four sampling methods can remove noises and generate realistic CT images. After averaging multiple Monte Carlo sampling results, excellent results can be obtained for all the four methods. In terms of sampling quality, the ODE method brings about inferior results while the other three methods produce comparable good results. However, in terms of sampling speed, the ODE method is significantly faster than the other methods while the PC method is the slowest. Practically, it is necessary to balance between sampling quality and sampling speed in an application-specific manner. Given the good sampling quality and a relative fast sampling speed of the EM method, we would recommend it as a good choice for SDE-based sampling.
When looking at the model uncertainty in Fig. 3(b) and standard deviation maps in Fig. 2, DDPM has a greater model uncertainty than the SDE-based sampling methods, we infer that the time-discrete training could be a major reason for the uncertainty of the DDPM model. We further assume that increasing the number of steps in the DDPM diffusion and reverse processes would allow DDPM to perform more like a time-continuous model with a reduced model uncertainty.
We have compared diffusion and score-matching results with CNN-based and GAN-based results. Among all the results, CNN results tend to be over-smoothed, which is partly due to the utilization of MSE as the objective function . GAN results have more details than the CNN results but are compromised by artifacts. These artifacts may come from a low robustness of the GAN model and a high tendency of hallucination. The diffusion and score-matching models, different from CNN and GAN models, have principled abilities to fit data distributions and generate high quality images. However, diffusion and score-matching models suffer from a well-known disadvantage: relying on a long Markov chain to generate results in a relatively slow speed.
As a future direction, we will keep exploring computational techniques to boost the sampling processes significantly. For that purpose, some algorithms were already proposed , but there is still a large gap between the speed of diffusion sampling and that of CNN and GAN inference.
Compared with the supervised approach, the unsupervised approach does not rely on paired images, and is more useful in medical imaging applications. However, the existing unsupervised diffusion and score-matching models have limitations. The constrains in the methods by Dharwal and Liu are too weak to reach specific structures or contents in generated images. On the other hand, the conditions in Song’s method are too strong, and it is hard to be implemented for CT-MRI synthesis as the similarity between CT and MRI data is not high. In future studies, we will explore opportunities to balance these conditional constrains for high quality CT-MRI synthesis.
In conclusion, we have adapted the emerging diffusion and score-matching models for image synthesis between CT and MRI. The four strategies, including DDPM, ODE, EM, and PC, have been adopted to sample CT images conditioned on an MRI image. The resultant CT images using different sampling strategies have been favorably compared with the results generated using conventional CNN and GAN models. The uncertainties of the diffusion and score-matching models have been quantified as well. Further investigation based on this paper is in progress.