FREA-Unet: Frequency-aware U-net for Modality Transfer

The goal of the proposed FREA-Unet model is to estimate a mapping $F_{MR\rightarrow PET}$ from the source image domain (MRI) to the target image domain (PET). The mapping F is learned from paired training data $S=\{(mr_{i},pet_{i})|mr_{i}\in MR,pet_{i}\in PET,i=1,2,...,N\}$, where N is the number of MRI-PET image pairs. Generating low/high frequency scales of the image separately with different loss functions can help the model to generate more realistic output with more preserved low/high frequency details. The proposed FREA-Unet model achieves this by explicitly generating low and high frequency PET scales in two different layers of the U-net’s decoder as pictured in Fig. 2.

There are two end-to-end-trainable attention modules for low/high frequency layers to weight different features based on their importance in PET image generation. Incorporating the attention mechanism into the layers responsible for generating low/high frequency scales helps the FREA-Unet model to focus more on the more relevant features during image translation leading to more realistic output images.

U-net architecture has two symmetrical paths: an encoding path and a decoding path to leverage both local and hierarchical information in order to estimate more accurate reconstruction of the encoded input. Similar to a regular CNN, the encoding path of the U-net contains convolutional layers to encode the context of the input image. On the other hand, the decoding path includes transposed convolutional layers to reconstruct the estimated image. The U-Net architecture has direct connections across its contracting path and expanding path. These skip connections combine high resolution features from the contracting path with the up-sampled features from the expanding path and use them as the inputs to the convolutional layers in the expanding path. This typically leads to improved resolution for the network output.

In our proposed FREA-Unet’s contracting path, we have six convolutional layers with 64, 128, 256, 512, 512, 512 filters, respectively. Each convolutional layer uses these filters to perform 2D convolutions on its input. The input of each layer is the output of its previous layer, except for the first layer in which the input is MRI image. The convolution outputs are followed by a nonlinear activation function: Rectified Linear Unit (ReLU), and batch normalization. Batch normalization allows using much higher learning rates and accelerates the network training. In the expanding path of FREA-Unet architecture, we have six transposed convolutional layers with 512, 1024, 1024, 512, 256, 128 filters, respectively, followed by ReLU and batch normalization. In addition, we use dropout in layers of the FREA-Unet as an effective technique for regularization and preventing the co-adaptation of neurons in the network.

In the proposed FREA-Unet architecture, we assign two different layers in the decoding path for low/high frequency image generation in order to generating low/high frequency PET images separately. Specifically, we consider the fourth layer of the decoding path for low frequency generation and the fifth layer of the decoding path for high frequency PET image generation. Low/high frequency scales are separated using a Gaussian filter for each image during the training phase so the model can optimize different scales separately. As it is illustrated in Fig. 2, low/high frequency scale outputs are combined after generation. As shown in our experimental results in Section 3, optimizing on different scales helps generate more realistic PET images with improvement in quantitative results.

Learning the end-to-end mapping function from MRI input and low/high frequency scales requires the estimation of network parameters achieved through minimizing the prediction error between the predicted images $F(MR;\theta)$ and the corresponding ground truth PET images. The objective for the mapping $F_{MR\rightarrow PET_{low}}$ can be defied as:

$$L_{low}(F)=||F(mr;\theta_{low})-pet_{low}||_{2}$$

(1)

Similarly, the objective for updating the mapping $F_{MR\rightarrow PET_{high}}$ in high frequency stream of the generator is defined as:

$$L_{high}(F)=||F(mr;\theta_{high})-pet_{high}||_{1}$$

(2)

Since $L_{1}$ norm leads to less blurry results in image generation tasks [5], we use $L_{1}$ norm in our high scale generation stream. Finally, by combining low/high frequency streams loss functions, the full objective of the FREA-Unet model can be expressed as:

$$L_{FREA-Unet}=L_{low}(F)+L_{high}(F)$$

(3)

In the proposed FREA-Unet model, low/high frequency scales are optimized separately. We incorporate attention mechanism into low/high scale layers to weight different features based on their importance in the output PET image. Specifically, the attention in the low/high frequency layers is defined as the compatibility scores [36] between feature maps extracted at frequency layers and the output from the second to the last layer in FREA-Unet’s decoding path showing the relevant features for generating the output image.

We define the compatibility score as result of the dot product between the feature map $f_{i}$ and the output of the second to last layer of the model that has the input image size, having only to pass through the final layer to produce the output PET image $O_{pet}$:

$$c_{i}=\langle f_{i},O_{pet}\rangle,i\in\{1....n\}$$

(4)

The extracted spatial attention scores are used in the low/high scale layers so that higher weights are given to the relevant regions, leading to more realistic output images.

The brain dataset was acquired from 30 subjects with both MRI and PET scans in the Alzheimer’s Disease Neuroimaging Initiative (ADNI) database (www. adni-info.org). ADNI provides a large database of studies with the goal of understanding the development and pathology of Alzheimer’s Disease. Subjects are diagnosed as cognitively normal (CN), significant memory concern (SMC), early mild cognitive impairment (EMCI), mild cognitive impairment (MCI), late mild cognitive impairment (LMCI) or having Alzheimer’s Disease (AD). All 30 PET subjects obtained from ADNI datasets were aligned to their corresponding T1- weighted MRI using coregistration so there is spatial correspondence MRI and PET slices in each subject. We cropped input MRI and PET images to the size of 256$\times$256 for training.

Extensive experiments were performed to compare FREA-Unet with current sate-of-the-art models for generating synthetic PET images from MRI inputs. Based on the same training and testing splits, the following models used previously for similar image generation tasks are included in our empirical evaluation and comparison.

• •

  U-net model: Following techniques developed by Han [11], we implemented and trained U-net model to compare against the FREA-Unet approach. The U-net model used in our empirical evaluation has no attention mechanism and generates only one frequency scale output. The architecture includes six convolutional layers in the encoding path and six transposed convolutional layers in the decoding path.

• •

  pix2pix: We implemented and trained pix2pix model developed by Isola [5], with five residual blocks in the generator. The discriminator is a regular CNN with five convolutional layers that classified the input image as real or synthetic.

• •

  cGANs with U-Net architecture (UcGAN): We implemented and trained conditional GAN (cGAN) model with U-net [17] architecture as its generator. We include comparison with UcGAN as FREA-Unet’s core architecture is also based on the U-net. The discriminator is a regular CNN with five convolutional layers.

To evaluate the performance, three-fold cross-validation [38] was used in our model training and testing. That is, out of the 30 cases, 20 cases are randomly selected for the training, and the remaining cases are used to test the trained model. We repeat the experiment three times and report the averaged outcome.

The weights in FREA-Unet were all initialized from a Gaussian distribution with parameter values of 0 and 0.02 for mean and standard deviation, respectively. The model is trained with ADAM optimizer with an initial learning rate of 0.0002 and with a batch size of 1. We trained the model for 200 epochs on a NVIDIA GTX 1080 Ti GPU. We used the same training setting for all models in our comparison.

We used three commonly-used quantitative measures to evaluate the prediction accuracy of the models between the generated PET images and the real PET. These include Mean Absolute Error (MAE) defined below:

$$MAE=\frac{\sum_{i=1}^{H}|realPET(i)-synPET(i)}{H}$$

(5)

where H is the number of body voxels. The second measure is the Structural Similarity Index (SSIM) defined as:

$$SSIM=\frac{(2\mu_{PET}\mu_{synPET}+C_{1})(2\sigma_{PETsynPET}+C_{2})}{(\mu_{PET}^{2}+\mu_{synPET}^{2}+C_{1})(\sigma_{PET}^{2}+\sigma_{synPET}^{2}+C_{2})}$$

(6)

where $\mu_{PET}$ denotes the mean value of PET image, $\mu_{synPET}$ denotes the mean value of generated PET image, $\sigma_{PET}^{2}$ is the variance of image PET, $\sigma_{synPET}^{2}$ is the variance of image synPET, and the parameters $C1=(k1Q)^{2}$ and $C2=(k2Q)^{2}$ are two variables to stabilize the division with weak denominators, where $k1=0.01$ and $k2=0.02$. The last measure is the Peak Signal-to-Noise Ratio (PSNR) defined as:

$$PSNR=10log_{10}(\frac{Q^{2}}{MSE})$$

(7)

where $Q$ is the maximal intensity value of PET and synPET images, and MSE is the mean squared error.

For the MAE metric, the lower value means better prediction results, that is, more realistic generated PET images. For both SSIM and PSNR, higher values indicate better prediction.

We first performed model ablation to evaluate the impact of each component of FREA-Unet. In Table 1, we report MAE, PSNR and SSIM for different configurations of our model. First, we removed the attention mechanism from the model and also optimized the model on the original input scale (without seperating low/high frequency scales). As a consequence, we also used the regular U-net loss. In this case, our model is reduced to the U-net architecture (U-net) with MAE: 52.56 $\pm$ 6.71, PSNR: 26.75 $\pm$ 1.96 and SSIM: 0.84 $\pm$ 0.05.

Next, we removed ”optimization on different scales” but kept attention mechanism to weight different features (FREA-Unet-wo-Freq). Our results show that this leads to better results (MAE: 51.10 $\pm$ 7.03, PSNR: 27.08 $\pm$ 1.33 and SSIM: 0.85 $\pm$ 0.04) than U-net, but worse than FREA-Unet. Clearly, separately optimizing low/high scales can help to achieve more accurate synPETS with preserved organ structures.

Further, we removed the attention mechanism from our model and optimized the model on low/high scales separately (FREA-Unet-wo-Att). This results to a higher MAE (48.19 $\pm$ 6.42) and lower PSNR (27.93 $\pm$ 1.84) and SSIM (0.86 $\pm$ 0.04) when compared with the full version of FREA-Unet. This ablation study showed that FREA-Unet works better when attention is employed to weight different features based on their importance in estimating PET images. Separately optimizing low/high scales using different weights obtained from attention modules helped to generate more realistic synthetic PETs with higher spatial resolution.

Fig. 3, shows generated synPET results using different models for different test cases. The images in the columns are the input MRI, real PET, and the generated PETs using the proposed FREA-Unet model, U-net [11], UcGAN [17], and pix2pix [5] respectively.

Qualitative comparisons for different methods show that the synthetic PETs generated by the proposed FREA-Unet are more accurate in estimating anatomical structures and have more preserved details compared to synPETs generated by other models. Since FREA-Unet generates low/high frequency scales separately, its generated results have higher spatial resolution, and have a successful translation in generating sharp regions including boundaries as shown in Fig. 3.

The average MAE, PSNR, and SSIM metrics computed based on the real and synthetic PETs for all test cases are listed in Table 2 for each of the aforementioned methods. The FREA-Unet achieved the best performance in synthetic PET generation compared to other models with the average MAE of 46.26$\pm$5.03 for all test cases.

In this paper, we propose a frequency-aware U-net model, FREA-Unet, for modality transformation in medical imaging. FREA-Unet optimizes on different scales separately and takes advantage of attention mechanism to weight different features in low/high scale generation. Separately optimizing low/high scales using different weights obtained from attention modules helped to generate more accurate synthetic PETs with preserved organ structures. We have applied this model to estimate PET images from their corresponding MR images on ADNI dataset where we are able to generate sharp synthetic PET images similar to real PET images. Our experiments show that the proposed FREA-Unet model can outperform other state-of-the-art methods in modality transfer of medical imaging.