A resource-efficient deep learning framework for low-dose brain PET image reconstruction and analysis

L-PET images are reconstructed through 5$\%$ undersampling of list-mode F-PET data. Both L-PET images and F-PET images are preprocessed by Statistical Parametric Mapping (https://www.fil.ion.ucl.ac.uk/spm/) for realignment and normalization. After preprocessing, the voxel size of these PET images is $1\times 1\times 1$ mm, and each 3D PET image is split into several closely stacked 2D PET image slices with $256\times 256$ pixels.

The proposed transGAN-SDAM consists of transGAN and SDAM, as illustrated in Fig. 1. Firstly, L-PET image slices are fed into transGAN, containing four sub-networks: transGAN Generator $G$, PatchGAN Discriminator $D$ [5], high feature encoder VGG16-Net and VGG19-Net [10], to generate the corresponding F-PET slices in a 2.5D manner. Different from previous GAN-based studies that directly adopting a uNet generator [11] or a ResNet generator [12, 13] for image generation, we design a transformer-encoded ResNet generator $G$ that can encode and optimize all image patches in whole-slice scale through the multi-headed self-attention mechanism before every generation process. Since all generation processes are pre-encoded, the transGAN converges more quickly and ensure the semantic stability of generated F-PET images. During training the proposed transGAN, the adversarial loss function $\mathcal{L}_{GAN}$ is defined as:

$$\mathcal{L}_{GAN}(D,G)=\\ -\mathbb{E}_{x,y}[(D(x,y)-1)^{2}]-\mathbb{E}_{x}[D(x,G(x))^{2}],$$

(1)

where $x$ denotes the L-PET slices, $G(x)$ denotes the F-PET image generated by the generator, $y$ is the corresponding ground-truth F-PET image and $\mathbb{E}$ represents the expectation. We also introduce the Charbonnier Loss [14] to punish the Euclidean difference between the synthesized F-PET images and the ground truth:

$$\mathcal{L}(G)=\mathbb{E}_{x,y}[\sqrt{\Arrowvert y-G(x)\Arrowvert^{2}+\epsilon^{2}}],$$

(2)

where $\epsilon$ denotes a small constant that is set to be $1\times 10^{-8}$. Considering the perceptual difference between the generated and the ground-truth F-PET images, inspired by [7, 12], we include both a VGG16-Net and a VGG19-Net that are pretrained on the ImageNet [15] to extract high feature representations of both $y$ and $G(x)$. The dual perceptual loss is defined as:

$$\begin{array}[]{r}\mathcal{L}_{perc}(G,V)=\mathbb{E}_{x,y}[\sqrt{\Arrowvert V_{1}(y)-V_{1}(G(x))\Arrowvert^{2}+\epsilon^{2}}]\\ +\mathbb{E}_{x,y}[\sqrt{\Arrowvert V_{2}(y)-V_{2}(G(x))\Arrowvert^{2}+\epsilon^{2}}],\end{array}$$

(3)

where $V_{1}(\cdot)$ and $V_{2}(\cdot)$ represent the feature maps extracted by VGG16-Net and VGG19-Net, respectively. Thus, the total loss of training transGAN is given as follows:

$$\mathcal{L}_{\text{trans}}(G,V,D)=\mathcal{L}_{GAN}(D,G)+\alpha\mathcal{L}(G)+\beta\mathcal{L}_{\text{perc}}\left(G,V\right),$$

(4)

where $\alpha$ and $\beta$ are the hyper-parameters that govern the trade off among $\mathcal{L}_{GAN}(D,G)$, $\mathcal{L}(G)$ and $\mathcal{L}_{perc}(G,V)$. The hyper-parameters $\alpha$ and $\beta$ are set to be 100 and 100, respectively.

To generate 3D F-PET images, we propose SDAM to better extract additional informative structural contents and semantic details with a deformable convolution layer [16, 17]. Different from existing methods [18, 19] that explored spatial information by neighboring slices in a pairwise manner, we employ a U-Net based network as the offset prediction network which is fed into a sequence of generated F-PET images, to predict all the deformable offset field $\Delta$ jointly for the target slice that needs be enhanced, as shown in Fig. 1. It demonstrates that the joint manner has a much lower computational cost than the pairwise scheme [18] since all the deformation offsets can be obtained in a single forward pass. Then, the offset fields are further adaptively fused into the sequence of generated F-PET images via a deformable convolutional layer to get fused feature maps:

$$\mathcal{F}(\mathbf{p})=\sum_{t={t_{0}-r}}^{t_{0}+r}\sum_{s=1}^{S^{2}}K_{t,s}\cdot C_{t}(\mathbf{p}+\mathbf{p}_{s}+\mathbf{\delta}_{(t,\mathbf{p}),s}),$$

(5)

where $\mathcal{F}(\cdot)$ denotes the derived feature map, $S^{2}$ represents the size of convolution kernel and $K_{t}\in\mathbb{R}^{S^{2}}$ denotes the kernel for $t$-th channel. $\mathbf{p}$ is a spatial position and $\mathbf{p}_{s}$ indicates the regular sampling offset. the deformation offsets $\mathbf{\delta}_{(t,\mathbf{p}),s}$ are position-specific, i.e., an individual $\mathbf{\delta}_{(t,\mathbf{p}),s}$ will be assigned to each convolution window centered at the spatial position $(t,\textbf{p})$. In particular, the preceding and succeeding $r$ slices are regarded as the reference, which together with the target slice $C_{t_{0}}$ input to SDAM. Thus, the offset field $\Delta\in\mathbb{R}^{(2r+1)\times 2S^{2}\times H\times W}$ for all the spatial positions in the sequence of F-PET slices:

$$\Delta=\mathcal{F}_{\theta_{os}}([C_{t_{0}-r},\cdots,C_{t_{0}},\cdots,C_{t_{0}+r}]),$$

(6)

where $[\cdot]$ denotes the concatenation of the generated F-PET slices, and $\mathcal{F}_{\theta_{os}}(\cdot)$ is the offset prediction network. $H$, $W$ is the height and width, respectively.

To fully explore the complementary information in the fused feature maps, we feed $\mathcal{F}$ through a ResNet-based reconstruction network to predict a restoration residue $\hat{C}_{t_{0}}^{HQ}$. The refined target slice $C_{t_{0}}^{HQ}$ can finally be obtained by adding the residue back to the target slice as:

$$C_{t_{0}}^{HQ}=\hat{C}_{t_{0}}^{HQ}+C_{t_{0}}.$$

(7)

Thus, the loss function $\mathcal{L}_{SDAM}$ is set to the sum of squared error between the refined target frame $C_{t_{0}}^{HQ}$ and the ground-truth $y$:

$$\mathcal{L}_{SDAM}={\Arrowvert y-C_{t_{0}}^{HQ}\Arrowvert}^{2}_{2}.$$

(8)

Finally, each subject’s all refined F-PET image slices are combined together to form the final 3D PET images for quantitative evaluations such as SUVR analysis.

As shown in Table 1, we can observe that transGAN-SDAM significantly outperforms all of the competing models (p-values $<$ 0.05 for all paired comparisons) with the highest PSNR (33.9$\pm$4.0), SSIM (0.950$\pm$0.04), and the lowest VSMD (0.043$\pm$0.039), indicating that the quality of generated F-PET images by transGAN-SDAM is closest to the ground-truth. It is worth noting that the transGAN-SDAM module substantially improves PSNR by at least 2.0 dB, SSIM by at least $1.1\%$, and decreases the VSMD by at least $24.6\%$ over these state-of-the-art approaches. We also find that SDAM can further refine the image quality to achieve higher PSNR, SSIM and lower VSMD when comparing transGAN and transGAN-SDAM. As shown in Fig. 2, we observe that transGAN and transGAN-SDAM can produce more accurate structural and metabolic details than compared approaches. The pseudo-color difference maps between the real and the generated F-PET images by different approaches are illustrated in Fig. 3, which shows that transGAN-SDAM achieves the smallest voxel-scale difference to the ground-truth F-PET images.

We further present the performance comparison between transGAN-SDAM and 3D CycleGAN [22] in Fig. 4. As show in Fig. 4 (a), we can observe that whole-brain SUVRs of transGAN-SDAM have smaller 95$\%$ limits of agreements (from -0.19 to 0.19), smaller 95$\%$ CI (from -0.018 to 0.018) and stronger correlation coefficient (0.831) than 3D CycleGAN (agreement from -0.26 to 0.26, CI from -0.024 to 0.024, correlation coefficient: 0.633). Meanwhile, 3D CycleGAN needs 90.32 M of parameters to train (Fig. 4 (b)) and $2.0027\times 10^{4}$ MB of GPU memory to implement (Fig. 4 (c)), which means much higher resource consumption than our transGAN-SDAM (15.77 M of parameters to train and $2.869\times 10^{3}$ MB of GPU memory to implement). It indicates that compared with 3D-based models, our proposal can satisfy specific needs, e.g., high response speed and limited GPU resources.