# Super Efficient Neural Network for Compression Artifacts Reduction and Super
Resolution

Wen Ma, Qiuwen Lou, Arman Kazemi, Julian Faraone, Tariq Afzal
Amazon Lab126
1100 Enterprise Way, Sunnyvale, CA 94089
<EMAIL_ADDRESS>

###### Abstract

Video quality can suffer from limited internet speed while being streamed by
users. Compression artifacts start to appear when the bitrate decreases to
match the available bandwidth. Existing algorithms either focus on removing
the compression artifacts at the same video resolution, or on upscaling the
video resolution but not removing the artifacts. Super resolution-only
approaches will amplify the artifacts along with the details by default. We
propose a lightweight convolutional neural network (CNN)-based algorithm which
simultaneously performs artifacts reduction and super resolution (ARSR) by
enhancing the feature extraction layers and designing a custom training
dataset. The output of this neural network is evaluated for test streams
compressed at low bitrates using variable bitrate (VBR) encoding. The output
video quality shows a 4-6 increase in video multi-method assessment fusion
(VMAF) score compared to traditional interpolation upscaling approaches such
as Lanczos or Bicubic.

##  1 Introduction

Video streaming has become a major entertainment activity in many households
around the world. With the recent popularity of AI/ML hardware accelerators,
neural network-based upscaler has been of interests to be integrated on edge
devices like 4K displays, laptops, mobile phones, TVs, to achieve real time
upscaling of video contents. One of the challenges is that when the internet
bandwidth/bitrate is low, the input videos suffer from compression artifacts
such as blocking, ringing, flickering artifacts, color changes, etc. Machine
learning and neural network-based algorithms have been widely studied to
address the problem of removing compression artifacts in videos, mainly based
on convolutional neural networks (CNNs) [18] [5]. Usually these models are
heavy weight and not feasible to run on resource constrained edge and mobile
devices, as they use multiple frames as the input to produce one output frame,
and often involves complicated algorithms such as optic flow and motion
compensation. In this paper, we introduce a hardware friendly network that
simultaneously performs artifacts reduction and super resolution (ARSR), by
combining super efficient neural network (SESR) [3], a lightweight CNN-based
solution to super resolution, with an artifacts reduction approach used in
ARCNN [6]. The network can produce output video quality close to a state-of-
the-art model, BasicVSR++, for video super resolution (VSR) and video
enhancement (VE) at a significantly reduced model size.

The major contributions of this paper are: 1) proposing a lightweight CNN-
based model with 22K parameters for ARSR. 2) using only one frame at a time to
produce one output upscaled/enhanced frame, 3) using over-parameterization on
both feature extraction and non-linear mapping layers to improve the inference
efficiency and picture quality, and 4) using video multi-method assessment
fusion (VMAF) to evaluate video quality in test datasets.

##  2 Background

During the past few years, the trend for image restoration and enhancement has
been pivoted from single image super resolution [20] to video super resolution
and real world video problems, especially on video enhancement and quality
improvement on compressed videos. BasicVSR++ [4], a winner of the NTIRE
challenge held in 2021, has been treated as a state-of-the-art model for video
super resolution and compressed video enhancement. The subsequent winners of
this challenge [22] are heavily adapted from this model. BasicVSR++ is based
on recurrent network and uses optical flow that uses second order grid
propagation and flow guided deformable alignment, both of which involve
complicated algorithms. In addition, the model has 7.3M parameters and cannot
be easily implemented in edge devices.

Another popular architecture to address the video super resolution and video
enhancement problems is the generative adversarial networks (GAN) [10]. The
advantage of using GAN is the ability to use little input data to generate
realistic outputs. A GAN architecture is made of two networks, a discriminator
network that distinguishes ground truth high resolution images from upscaled
ones, and a generator network that produces images that gradually look like
ground truth images. There is a notable work called SUPERVEGAN [2] that uses
GAN network for video super resolution. SUPERVEGAN combines dynamic upsampling
filters (DUF) and GAN to perform artifacts removal and super resolution at the
same time. However, it is using 5 low resolution input frames to generate a
single high resolution frame. In addition, the model is developed initially as
a software approach to address the VSR and VE problems and is not suitable for
hardware implementations.

Super Resolution algorithms based on simple stacking of convolution layers is
favored due to the potential simplicity of the algorithms and models. Several
lightweight models have been proposed, such as FSRCNN [8], ABRL [9], LapSRN
[14], CARN-M [1], etc. The number of weight parameters can be reduced to less
than 30K in some of these models and the amount of MACs can be reduced to less
100 billion when upscaling images from 360p to 720p resolution. A notable
solution in this category is the super efficient super solution (SESR) network
[3], which uses over-parameterization and collapsible linear blocks to reduce
the inference complexity, making it suitable for hardware inference
operations. However, SESR only performs super resolution but not artifacts
reduction. Recently, SRCNN [7] model architecture has been modified to ARCNN
[6] to deal with compression artifacts removal. The number of convolutional
layers in ARCNN was expanded for feature extraction and feature enhancement.
The non-linear mapping and reconstruction layers remained the same. Inspired
by [6], we came up with the idea of extending SESR for artifacts reduction by
adopting similar techniques used for ARCNN. Details will be elaborated in
section 3.

Video multi-method assessment fusion, or VMAF [15] was developed for
perceptual video quality assessment. The other video quality metrics such as
PSNR or SSIM that use pixel-based calculations may fail to capture human
perception accurately [16]. VMAF, on the other hand, combines human vision
modeling with support vector machines to produce a score that represents human
perception. The final metric score is calculated using the weighted scores
from several elementary metrics. The elementary metrics scores were obtained
using opinion scores through subjective experiments. We adopted this method as
it is open sourced and can be directly used along with FFmpeg commands. For
measuring video quality shown on 4K TVs, we used the VMAF model for 4K
resolution while disabling enhancement gain mode (NEG).

Figure 1: Model architecture of the ARSR network. The low resolution (LR)
frame goes through $N$ Conv2d layers for feature extraction, $M$ Conv2d layers
for non-linear mapping, a Conv2d layer to match the upscaling factor, and a
final depth-to-space layer to upscale the frame to super resolution (SR). For
later experiments, unless specified, $N=3$ and $M=11$ are used.

##  3 Super Efficient Neural Network for Artifacts Reduction and Super
Resolution

Our neural network architecture was adapted from SESR [3] which does not
perform artifacts reduction and focuses on single image super resolution.
Various modifications of the neural network have been tested to optimize the
performance, such as architecture change, training dataset modifications, loss
function variations, different chroma channel upscaling techniques. Also the
network was studied to support different upscaling factors such as $\times 2$,
$\times 3$ and $\times 4$. To make the network hardware friendly, techniques
such as grouped convolutions, quantization aware training, power-of-two
scaling were tested on video enhancement tasks.

###  3.1 Overall ARSR Neural Network Architecture

The architecture of the CNN-based ARSR Network is shown in Fig. 1. The network
only processes the Y channel component of single images. There is a global
residual layer feeding the input image (Y channel only) to the depth-to-space
layer at the end. On the other side, the input image goes through $N$ layers
of feature extraction. A second residual layer starts from the beginning of
non-linear mapping layers and ends at the end of the non-linear mapping
layers. The number of non-linear mapping layers ($M$ in Fig. 1) can range from
5 to 11 depending on the image quality or hardware cost requirements. The
depth-to-space layer at the end performs the actual upscaling of images,
therefore relaxing the memory requirements because the intermediate outputs of
previous layers are at the low resolution (LR) image sizes.

Figure 2: During training, the Conv2d operation is expanded into two
sequential Conv2d operations. This is an example where the number of input and
output channels are both 16 while the number of internal expanded channel is
256.

The Conv2d layers all use over-parameterization techniques where a normal
convolution operation with $n$ input channels, $m$ output channels, and
$f\times f$ kernel sizes will be expanded to two convolution operations where
the first one has $n$ input channels, $p$ output channels, with kernel size of
$f\times f$, while the second one has $p$ input channels, $m$ output channels,
with kernel size of $1\times 1$. Usually the value of $p$ is much larger than
$m$ and $n$. During training, the weight updates are done on the expanded
network, with $npff+pm$ weights. During the inference, only the $nmff$
collapsed weights are used, leading to more much efficient inference
operations. The over-parameterization scheme is shown in Fig. 2.

###  3.2 Training Dataset Generation

In stead of using the traditional bicubic downscaled images as the low
resolution samples in the training dataset, we compressed and downscaled the
original videos from the Vimeo dataset [21] using H.265 codec at a low bitrate
- 50 kbps. We used extracted frames of the 4K resolution videos from the Vimeo
dataset as the high resolution samples, and the extracted frames of the
compressed and downscaled videos as the low resolution samples. Examples of
the low resolution and high resolution video frames are shown in Fig. 3. We
also tried another method of injecting JPEG compression artifacts into the
traditional bicubic downscaled low resolution samples but this method tends to
smooth out everything in the output image. Our adopted method can smooth out
area where there are artifacts and still keep the high frequency details.

(a) Low resolution (LR) image for training, compressed at 50 kbps using H.265
codec.

(b) High resolution (HR) image for training.

Figure 3: Examples from the Vimeo dataset for training.

(a) Input 540p frame with compression artifacts.

(b) Output 4K frame trained with MAE loss.

(c) Output 4K frame trained with MSE loss.

(d) Output 4K frame trained with Huber loss.

Figure 4: Comparison between different loss functions during training.

(a) Input 540p frame with compression artifacts.

(b) Upscaled 4K frame with 1 layer of feature extraction.

(c) Upscaled 4K frame with 2 layers of feature extraction.

(d) Upscaled 4K frame with 3 layers of feature extraction.

Figure 5: Using different number of layers of feature extraction to reduce
artifacts.

The dataset generated is based on single frame, meaning that one input image
frame is used to generate one output frame. It is not uncommon that multiple
input frames are used in CNN to generate one output frame [12] for video super
resolution tasks. We explored the multi-frame approach by using three input
frames as the input and the high resolution version of the middle frame as the
output. The result looks almost the same as the single frame approach.

(a) Using Bilinear interpolation for chroma channels.

(b) Using Bicubic interpolation for chroma channels.

(c) Using Nearest Neighbor interpolation for chroma channels.

Figure 6: Using different chroma channels upscaling methods. Input Resolution |  Output Resolution |  Upscale Factor |  Number of Output Channels in the Last Conv2d Layer |  Use of Lanczos upscaler or not
---|---|---|---|---
$960\times 540$ |  $3840\times 2160$ |  4 |  16 |  No
$1280\times 720$ |  $3840\times 2160$ |  3 |  9 |  No
$1920\times 1080$ |  $3840\times 2160$ |  2 |  4 |  No
$960\times 540$ |  $2560\times 1440$ |  8/3 |  4 |  Yes
$1280\times 720$ |  $2560\times 1440$ |  2 |  4 |  No
$2560\times 1440$ |  $3840\times 2160$ |  3/2 |  N/A |  Yes
Table 1: ARSR network can support different upscaling factors combining
lanczos upscaler when needed.

###  3.3 Loss Function Exploration

Different kinds of loss functions during training have been explored,
including the mean average error (MAE), mean squared error (MSE), and Huber
loss. We found that MAE performs the best at artifacts reduction while Huber
loss and MSE preserves slightly more details at the cost of not removing
enough artifacts. A comparison between using different loss functions during
training is shown in Fig. 4.

###  3.4 Expanding Feature Extraction Layers for Artifacts Reduction

Inspired by [6], we increased the number of feature extraction layers ($N$ in
Fig. 1) in the beginning of the network from one layer to multiple layers to
perform better artifacts reduction. More feature extraction layers can extract
more features from the input or training images to construct better output
images and produce better training results. For three layers of feature
extraction, we adopted the $7\times 7$ kernel size in the first layer,
followed by $5\times 5$ and $3\times 3$ kernel size in the second and third
layers. For two layers of feature extraction, we adopted the $7\times 7$
kernel size in the first layer, followed by $5\times 5$ kernel size in the
second layer. The number of input and output channels are both 16. We found
that using two or three layers of feature extraction is sufficient to remove
the artifacts, as shown in Fig. 5.

###  3.5 Chorma Channels Upscaling Methods

While the ARSR neural network is used to upscale Luma channel (Y channel)
only, traditional upscaling method is used to upscale the chroma channels (Cb,
Cr channels). From [7], we know that using neural network-based upscaler on
luma channel only is sufficient to generate competitive super resolution
results. We explored using bicubic, bilinear and nearest neighbor methods for
chroma upscaling, and found that the bilinear method produces the smoothest
edges, while nearest neighbor method tend to add some artifacts near the edges
(Fig. 6). The nearest neighbor method is the easiest to implement in hardware
as it has the lowest computation complexity, whereas the bicubic method is the
most difficult. We have decided to use bilinear upscaling for chroma channels
because it can remove the jaggedness on edges and produce smooth-looking
outputs.

(a) Upscaled 4K frame by Lanczos interpolation.

(b) Upscaled 4K frame by ARSR quantized network.

(c) Upscaled 4K frame by ARSR FP32 network.

(d) Upscaled 4K frame by BasicVSR++.

Figure 7: Comparison between different upscaling methods.

(a) Upscaled 4K frame by Lanczos interpolation.

(b) Upscaled 4K frame by ARSR quantized network.

(c) Upscaled 4K frame by ARSR FP32 network.

(d) Upscaled 4K frame by BasicVSR++.

Figure 8: One example showing too much artifacts reduction can lead to loss of details. Upscaling Method |  Average VMAF of 8 test videos |  Number of Parameters
---|---|---
ARSR FP32 |  52.89 |  41.2K
ARSR quantized |  51.64 |  22.2K
Lanczos |  47.55 |  N/A
Bicubic |  45.37 |  N/A
BasicVSR++ |  55.68 |  7.3M
Table 2: VMAF score comparison when upscaling 540p compressed videos to 4K
videos. ARSR network is producing higher VMAF scores than Lanzcos and Bicubic
upscalers.

We have also tried to modify the neural network to be 3-channel models, either
based on YCbCr or RGB. We found that 3-channel models can produce slightly
better output image quality but there is color difference between the expected
output and generated output. Therefore we didn’t pursue this method further.

###  3.6 Supporting Different Upscaling Scenarios

The ARSR neural network can support $\times 2$, $\times 3$, and $\times 4$
upscaling factors by changing the number of output channels in the last Conv2d
layer shown in Fig. 1, before the depth-to-space layer. If the upscaling
factor is $n$, the number of output channels would be $n^{2}$. So a larger
integer upscaling factor would require slightly more parameters in the neural
network. Tab. 1 shows different upscaling scenarios for different input video
resolutions and output video resolutions. For upscaling scenarios where the
upscaling factor is not an integer, we upscale using the biggest integer that
will result in an output resolution lower than the desired output resolution
and then use the lanczos interpolation method to upscale to the desired output
resolution.

###  3.7 Techniques to Improve the Efficiency of ARSR Neural Network

To reduce the number of parameters and computations in the network, we propose
to use the grouped convolutions [13] to divide the 16 input and 16 output
channels of the non-linear mapping layers into $n$ groups. The value of $n$
can be 8, 4 or 2. Normal Conv2d operations are performed within each group of
input channels and output channels. The outputs from each group are finally
concatenated into the final output. The number of model parameters and
computations of the non-linear mapping layers is reduced by $n$.

Other techniques we adopted to make the network suitable for hardware
implementation include quantization-aware training [11] and power-of-two
scaling factor during quantization [17]. Power-of-two scaling during
quantization can convert the multiplication operation during
quantization/dequantization into bit-shift operations to simplify the
computation. However, using power-of-two scale factors can lead to lower
accuracy due to the rounding errors so we didn’t integrate this technique into
our quantized model.

##  4 Experiment Results of the ARSR Neural Network

Instead of using the popular PSNR or SSIM metrics to evaluate picture quality,
we use the VMAF score to measure quality as it can better evaluate the degree
of artifacts reduction in addition to detail enhancement. The model of
VMAF_4K_NEG is used because it measures the audience’s response when they are
viewing the video on a 4K TV. We selected 8 videos from the inter4K dataset
[19] and generated a testing dataset for ARSR evaluation. We compressed the
video using variable bitrate (VBR) encoding at 500 kbps and then upscale by
ARSR network or other comparable methods. The baseline we chose are lanczos,
which is the most complicated traditional interpolation method, and bicubic
upscaling approaches.

We demonstrated the results of ARSR network by showing the output of a
quantized $\times 4$ network that uses group of 4 convolutions and 12 bit on
weights and activations. The test inputs are 540p compressed videos. The test
outputs are upscaled and enhanced 4K videos. In comparison, the outputs of the
ARSR FP32 (Floating Point 32-bit) network, the state-of-the-art BasicVSR++
network, and baseline Lanczos upscaler are shown as well (Fig. 7). We see that
the Lanczos upscaler does not remove compression artifacts or mosquito noise,
while the BasicVSR++ performs the best at removing artifacts. However,
BasicVSR++ can sometimes remove too much artifacts that it smooths out the
details, as shown in Fig. 8. The outputs of the quantized ARSR network is
slightly worse than the FP32 network, but still better than the Lanczos
upscaled outputs (Fig. 7).

To evaluate the video quality quantitatively, we measured the VMAF scores of
the 8 upscaled output videos and calculated the average. The result is shown
in Tab. 2. BasicVSR++ has the highest VMAF score, which indicates the best
video quality, followed by ARSR FP32 network and ARSR quantized network. On
the other hand, BasicVSR++ also uses the most complex model with 7.3M
parameters, whereas ARSR FP32 has 41.2K parameters and ARSR quantized only has
22.2K parameters. The Lanczos and Bicubic methods have the lowest VMAF scores.
We believe this table has shown that our ARSR network and training method is
effective in producing noticeable improvement in video quality.

##  5 Conclusion

We have developed a CNN-based neural network for simultaneous artifacts
removal and super resolution (ARSR). It is lightweight, based on single frame,
and easy to implement in hardware. Based on a selected test dataset, the
performance of ARSR is about 4-6 higher in the VMAF score than lanczos or
bicubic methods. The state-of-the-art model, BasicVSR++ which has 7.3M
parameters, is not feasible to run on resource constrained edge device
hardware. This is the first known hardware friendly network that performs
artifacts reduction and super resolution at the same time.

Some of the future work can include recovering the picture quality loss from
the quantized model, making different training dataset for different input
video bitrate/resolution scenarios, and further reducing the network size by
decreasing the number of non-linear mapping layers, e.g. $M=5$.

## References

  * [1] Namhyuk Ahn, Byungkon Kang, and Kyung-Ah Sohn.  Fast, accurate, and lightweight super-resolution with cascading residual network.  In Computer Vision – ECCV 2018: 15th European Conference, Munich, Germany, September 8-14, 2018, Proceedings, Part X, page 256–272, Berlin, Heidelberg, 2018. Springer-Verlag.
  * [2] Silviu S. Andrei, Nataliya Shapovalova, and Walterio Mayol-Cuevas.  Supervegan: Super resolution video enhancement gan for perceptually improving low bitrate streams.  In IEEE Access, volume 9, pages 91160–91174, 2021.
  * [3] Kartikeya Bhardwaj, Milos Milosavljevic, Liam O' Neil, Dibakar Gope, Ramon Matas, Alex Chalfin, Naveen Suda, Lingchuan Meng, and Danny Loh.  Collapsible linear blocks for super-efficient super resolution.  In D. Marculescu, Y. Chi, and C. Wu, editors, Proceedings of Machine Learning and Systems, volume 4, pages 529–547, 2022.
  * [4] Kelvin C.K. Chan, Shangchen Zhou, Xiangyu Xu, and Chen Change Loy.  Basicvsr++: Improving video super-resolution with enhanced propagation and alignment.  In 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 5962–5971, 2022.
  * [5] Wei-Gang Chen, Runyi Yu, and Xun Wang.  Neural network-based video compression artifact reduction using temporal correlation and sparsity prior predictions.  IEEE Access, 8:162479–162490, 2020.
  * [6] Chao Dong, Yubin Deng, Chen Change Loy, and Xiaoou Tang.  Compression artifacts reduction by a deep convolutional network.  In 2015 IEEE International Conference on Computer Vision (ICCV), pages 576–584, 2015.
  * [7] Chao Dong, Chen Change Loy, Kaiming He, and Xiaoou Tang.  Image super-resolution using deep convolutional networks.  IEEE Trans. Pattern Anal. Mach. Intell., 38(2):295–307, feb 2016.
  * [8] Chao Dong, Chen Change Loy, and Xiaoou Tang.  Accelerating the super-resolution convolutional neural network.  In Bastian Leibe, Jiri Matas, Nicu Sebe, and Max Welling, editors, Computer Vision – ECCV 2016, pages 391–407, Cham, 2016. Springer International Publishing.
  * [9] Z. Du, J. Liu, J. Tang, and G. Wu.  Anchor-based plain net for mobile image super-resolution.  In 2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops (CVPRW), pages 2494–2502, Los Alamitos, CA, USA, jun 2021. IEEE Computer Society.
  * [10] Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio.  Generative adversarial nets.  In Advances in neural information processing systems, pages 2672–2680, 2014.
  * [11] Benoit Jacob, Skirmantas Kligys, Bo Chen, Menglong Zhu, Matthew Tang, Andrew Howard, Hartwig Adam, and Dmitry Kalenichenko.  Quantization and training of neural networks for efficient integer-arithmetic-only inference.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 2704–2713, 2018.
  * [12] Armin Kappeler, Seunghwan Yoo, Qiqin Dai, and Aggelos K Katsaggelos.  Video super-resolution with convolutional neural networks.  IEEE transactions on computational imaging, 2(2):109–122, 2016.
  * [13] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton.  Imagenet classification with deep convolutional neural networks.  In F. Pereira, C.J. Burges, L. Bottou, and K.Q. Weinberger, editors, Advances in Neural Information Processing Systems, volume 25. Curran Associates, Inc., 2012.
  * [14] Wei-Sheng Lai, Jia-Bin Huang, Narendra Ahuja, and Ming-Hsuan Yang.  Deep laplacian pyramid networks for fast and accurate super-resolution.  In 2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 5835–5843, 2017.
  * [15] Zhi Li, Anne Aaron, Ioannis Katsavounidis, Anush Moorthy, Megha Manohara, et al.  Toward a practical perceptual video quality metric.  The Netflix Tech Blog, 6(2):2, 2016.
  * [16] Juan Carlos Mier, Eddie Huang, Hossein Talebi, Feng Yang, and Peyman Milanfar.  Deep perceptual image quality assessment for compression.  In 2021 IEEE International Conference on Image Processing (ICIP), pages 1484–1488. IEEE, 2021.
  * [17] Dominika Przewlocka-Rus, Syed Shakib Sarwar, H Ekin Sumbul, Yuecheng Li, and Barbara De Salvo.  Power-of-two quantization for low bitwidth and hardware compliant neural networks.  arXiv preprint arXiv:2203.05025, 2022.
  * [18] Jae Woong Soh, Jaewoo Park, Yoonsik Kim, Byeongyong Ahn, Hyun-Seung Lee, Young-Su Moon, and Nam Ik Cho.  Reduction of video compression artifacts based on deep temporal networks.  IEEE Access, 6:63094–63106, 2018.
  * [19] Alexandros Stergiou and Ronald Poppe.  Adapool: Exponential adaptive pooling for information-retaining downsampling.  IEEE Transactions on Image Processing, 32:251–266, 2022.
  * [20] Radu Timofte, Shuhang Gu, Jiqing Wu, and Luc Van Gool.  Ntire 2018 challenge on single image super-resolution: Methods and results.  In Proceedings of the IEEE conference on computer vision and pattern recognition workshops, pages 852–863, 2018.
  * [21] Tianfan Xue, Baian Chen, Jiajun Wu, Donglai Wei, and William T Freeman.  Video enhancement with task-oriented flow.  International Journal of Computer Vision, 127:1106–1125, 2019.
  * [22] Ren Yang, Radu Timofte, Meisong Zheng, Qunliang Xing, Minglang Qiao, Mai Xu, Lai Jiang, Huaida Liu, Ying Chen, Youcheng Ben, et al.  Ntire 2022 challenge on super-resolution and quality enhancement of compressed video: Dataset, methods and results.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 1221–1238, 2022.


# Longitudinal Prediction of Postnatal Brain Magnetic Resonance Images via a
Metamorphic Generative Adversarial Network

Yunzhi Huang  Sahar Ahmad  Luyi Han  Shuai Wang  Zhengwang Wu  Weili Lin  Gang
Li  Li Wang  Pew-Thian Yap<EMAIL_ADDRESS>School of Automation, Nanjing
University of Information Science and Technology, Nanjing 210044, China
Department of Radiology and Biomedical Research Imaging Center (BRIC),
University of North Carolina, Chapel Hill, USA  Department of Radiology and
Nuclear Medicine, Radboud University Medical Center, Geert Grooteplein 10,
6525 GA, Nijmegen, The Netherlands  Department of Computer Science, Shandong
University (Weihai), China

###### Abstract

Missing scans are inevitable in longitudinal studies due to either subject
dropouts or failed scans. In this paper, we propose a deep learning framework
to predict missing scans from acquired scans, catering to longitudinal infant
studies. Prediction of infant brain MRI is challenging owing to the rapid
contrast and structural changes particularly during the first year of life. We
introduce a trustworthy metamorphic generative adversarial network (MGAN) for
translating infant brain MRI from one time-point to another. MGAN has three
key features: (i) Image translation leveraging spatial and frequency
information for detail-preserving mapping; (ii) Quality-guided learning
strategy that focuses attention on challenging regions. (iii) Multi-scale
hybrid loss function that improves translation of tissue contrast and
structural details. Experimental results indicate that MGAN outperforms
existing GANs by accurately predicting both contrast and anatomical details.

###### keywords:

Infant brain MRI, Longitudinal prediction, Metamorphic GAN

††journal: Journal of Pattern Recognition

##  1 Introduction

Brain MRI is commonly used to investigate normative and aberrant brain
evolution through infancy [1]. To precisely chart brain growth trajectories,
_temporally dense_ longitudinal datasets are often required but are difficult
to acquire. Moreover, infant studies often involve incomplete longitudinal
datasets, given the unique challenges associated with infant MRI acquisition.
The missing data at different time points can be due to subject dropouts or
failed scans owing to excessive motion, insufficient coverage, or imaging
artifacts [2].

Longitudinal prediction of infant brain scans is challenging as brain MRI
contrasts change rapidly through the first year of life. The brain volume
doubles to about $65\%$ of the adult brain by the end of the first year [3].
The gray matter (GM) follows a faster growth trajectory ($108\%-149\%$
increase) compared to white matter (WM; $11\%$ increase) [4]. The rapid brain
evolution is characterized by both structural and contrast variations [5, 6].
As shown in Figure 1, the WM appears to be darker than the GM during the
neonatal phase as the brain is going through myelination, and by sixth month,
WM and GM are almost indistinguishable due to the poor tissue contrast.

Figure 1: Appearance and structural changes at two time points during the
first year of life. Wavelet decomposition for capturing structural details.

###  1.1 Related Work

The longitudinal prediction of infant brain MRI can be formulated as an image-
to-image translation task — mapping images from a source time point to a
target time point [7]. Several studies in the field of computer vision [7, 8,
9, 10] have shown that generative adversarial networks (GANs) [11] yield
superior performance in translating images from a domain to another. In the
field of medical image analysis, [12] introduced an auto-context GAN to
progressively refine MRI-to-CT synthesis. In their follow-up study, [13]
incorporated difficulty-aware attention mechanism to improve predictions in
challenging regions. Similarly, [14] introduced self-attention to encourage
the transformation of a foreground object while retaining the background.
Medical image-to-image translation network (MedGAN) [15] uses a pre-trained
classification network as feature extractor to match textures and structural
details of synthetic and target CT images. All the aforementioned methods for
cross-modality synthesis focus on appearance changes and neglect morphological
changes. The longitudinal prediction of infant MR brain images, however,
requires dealing with fast-paced structural and appearance changes.

To promote structural consistency in cross-modality synthesis, several recent
approaches incorporate segmentation similarity as a learning constraint [16,
17]. However, tissue segmentation of infant brain MRI is challenging due to
the overlap of GM and WM intensity distributions (Fig. 1). Several approaches
attempted to ensure structural consistency without relying on tissue maps.
[18] employed gradient differences in a loss function to improve the
prediction of boundaries. [19] incorporated gradient correlation differences
in a structure-consistency loss to improve edge alignment in MRI-to-CT
synthesis. Although successful, the gradient-based constraint introduces noise
and fail to capture sufficient boundary information in images with low
contrast. [20] incorporated a patch-based self-similarity loss by comparing
each patch with all its neighbors in a pre-defined non-local region to ensure
structural consistency. However, the search for corresponding non-local
regions is computationally expensive.

###  1.2 Contributions

In this paper, we employ CycleGAN [8], a cycle consistent generation
framework, to simultaneously learn structural and appearance changes between
two time points. Major contributions of our work are summarized below:

  1. (i)

We propose a trustworthy adversarial learning metamorphosis framework that
accounts for both the appearance and structural changes in infant brain MRI.

  2. (ii)

We use a spatial-frequency transfer block equipped with wavelet decomposition
to transform features from multiple frequency bands to learn the structural
changes.

  3. (iii)

We employ a quality guidance strategy to incorporate a quality-driven loss
function to improve predictions in challenging regions.

  4. (iv)

We devise a multi-scale hybrid loss function to improve the matching of both
the textural details and the anatomical edges between the predicted image and
the desired target image. The discriminator network is evoked at multiple
resolutions via deep-supervision, thus allowing accurate prediction of
anatomical structures through adversarial learning.

The rest of the paper is organized as follows: Section 2 details the proposed
method. Section 3.2 describes the dataset used for evaluation and presents the
experimental results. Section 4 provides additional discussion and concludes
the paper.

##  2 Methods

In this work, we implement a framework for prediction of metamorphic changes
using a GAN. Details of our method are described next.

###  2.1 Network Architecture

We propose a metamorphic GAN (MGAN) to predict the infant brain MR image
scanned at time point $t_{b}$ from a time point $t_{a}$. Without loss of
generality, we assume that $t_{b}>t_{a}$. Our network architecture, shown in
Fig. 2, is cycle-consistent and learns a reversible translation between the
two time-points. It consists of (i) a forward path for earlier-to-later time-
point image prediction and (ii) a backward path for later-to-earlier time-
point image prediction. The two generators $G_{a}$ and $G_{b}$ and their
corresponding discriminators $D_{a}$ and $D_{b}$ follow an encoder-decoder
architecture. Both the generators incorporate a spatial-frequency transfer
(SFT) block to transform the appearance and structural features via multiple
branches detailed in Fig. 3. The two discriminators estimate voxel-level
uncertainty maps, enabling the corresponding generators to focus on
challenging regions. We will describe the components of our network in the
subsequent sections.

Figure 2: Overview of the metamorphic GAN.

####  2.1.1 Metamorphic Generator

The metamorphic generator (Fig. 3) takes a 3D patch of size $64\times 64\times
64$ as input and predicts a 3D patch. The generator consists of an encoder,
SFT block, and a decoder.

Figure 3: Network architecture of the metamorphic generator.

##### Encoder

The encoding path consists of two convolution blocks, each with a $3\times
3\times 3$ convolution layer, followed by $3$D instance normalization (IN)
[21] and a rectified linear unit (ReLU) [22]. For downsampling, we use
convolution with a stride of $2$ instead of pooling to avoid potential
information loss. We keep a $1$-stride convolution in the first stage of the
encoder to retain details, and use a $2$-stride convolution in the second
stage. The resulting numbers of feature maps in the two-stage encoder are $64$
and $32$.

##### Spatial-frequency transfer block

Longitudinal prediction requires translating both contrast and structure
between two time points. We propose to embed a spatial-frequency transfer
block in between enocder-decoder to extract the spatial and frequency domain
information of feature maps. The SFT block is divided into two branches: (i)
frequency transform branch, and (ii) spatial transform branch. The frequency
transform branch is equipped with discrete wavelet transform (DWT) that takes
into account the low frequency tissue contrast and high frequency structural
details. The DWT layer decomposes the feature map into low frequency
approximation and high frequency details along three dimensions, resulting in
eight subvolumes: $LLL,LLH,LHL,LHH,HLL,HLH,HHL$ and $HHH$. This decomposition
allows more effective transfer of spatial-frequency details.

Given the $i$-th channel feature map $f^{i}$ of size $(s_{x}\times s_{y}\times
s_{z})$, the decomposed feature map $f^{i}_{j}$ at frequency band $j$ is
obtained by convolving $f^{i}$ with wavelet filter $w_{j}$:

$f_{j}^{i}=f^{i}\circledast w_{j}.$ (1)

The wavelet filters for each frequency band are calculated by DWT
decomposition and are preset in the convolution layer. Correspondingly, the
feature maps are reconstructed in the decode path via inverse discrete wavelet
transform (IDWT) layer. We show the representative feature maps from the DWT
and IDWT layers in Fig. 4. The DWT layer is akin to pooling layer as the DWT
decomposition halves the size of the input feature maps. The IDWT layer
corresponds to the deconvolution operation with the fixed weights obtained via
wavelet filters. There is also an intermediate _transfer_ operation between
the DWT and IDWT layer. This transfer operation is realized through $9$
residual blocks [23]; the input of each block is processed by two $3\times
3\times 3$ convolution layers with $64$ channels followed by IN and ReLU for
activation. A shortcut connection is added between the input and the output of
every residual convolution block. Residual transfer learning simplifies
feature generation and transfer from a source domain to a target domain.

Figure 4: Feature maps obtained from the DWT and IDWT layers.

The second branch in the SFT block — spatial transform branch — is integrated
to compensate for the information truncated by the wavelets. It is implemented
using a convolution layer with kernel size $3\times 3\times 3$ and stride of
$2$ to downsample the feature maps in spatial domain. These downsampled
feature maps undergo transfer operation and are later upsampled by strided
deconvolution layer with kernel size $3\times 3\times 3$.

Each branch in the SFT block is trained independently without weight sharing.
The feature maps from both the frequency and spatial transform branches are
concatenated using a $3\times 3\times 3$ convolution layer and stride of $1$,
followed by IN and ReLU operation; capturing both the contrast and structural
information for translating from the source domain to the target domain.

##### Decoder

Deep supervision [24] is leveraged in the decoding path to strengthen the
gradient flow and encourage learning useful representations at multiple
scales. The feature maps are upsampled by a $2$-stride deconvolution layer and
are then convolved with a $3\times 3\times 3$ kernel to get the predicted
output.

####  2.1.2 Uncertainty Quantization

The uncertainty associated with the prediction stems from two aspects:
epistemic uncertainty (model uncertainty) and aleatoric uncertainty (data
uncertainty) [25, 26]. As shown in Fig. 3, two Monte-Carlo (MC) dropout layers
are incorporated in our generator to estimate the epistemic uncertainty. MC
dropout regularizes the network weights as Bernoulli distributions for
variational Bayesian inference [27, 28]. Note, MC dropout is only enabled
during inference. A set of predictions
$\\{\hat{y}_{1},\hat{y}_{2},\dots,\hat{y}_{N}\\}$ are sampled from the
distribution $p(\hat{y}|I,\mathbf{w}_{n})$ via $N$ stochastic inferences using
the metamorphic generator. The epistemic uncertainty is estimated as the
variance over the predictions:

$\mathcal{U}_{e}=\sqrt{\dfrac{\sum_{n=1}^{N}(\hat{y}_{n}-\overline{y})^{2}}{N}},$
(2)

where $\hat{y}$ denotes to the prediction by feeding the generator $G$ an
input image $I$, $\mathbf{w}_{n}$ represents the generator weights after the
$n$-th dropout, $N$ refers to the number of prediction instances, and
$\overline{y}$ is the mean of the predictions.

The aleatoric uncertainty is typically measured with the test-time
augmentation technique [29, 30]. During inference, we perturb the input data
with spatial transformations (flip and rotation) and random noise. Similar to
the estimation of the epistemic uncertainty, we sample a set of predictions
$\\{\hat{y}_{1},\hat{y}_{2},\dots,\hat{y}_{N}\\}$ from the distribution
$p(\hat{y}|I,S))$ and estimate the aleatoric uncertainty as the variance over
the predictions:

$\mathcal{U}_{a}=\sqrt{\dfrac{\sum_{n=1}^{N}(S^{-1}(\hat{y}(S(x+rn))-\overline{y})^{2}}{N}},$
(3)

where $S$ represents the spatial transformation, $S^{-1}$ corresponds to the
inverse transformation, and $rn$ corresponds to random noise.

####  2.1.3 Multi-scale discriminator

The discriminator in MGAN has a U-shaped architecture, as shown in Fig. 5, to
locally distinguish the predicted images from real images. It takes as input a
$64\times 64\times 64$ image patch and outputs the quality probability map for
the given $3$D patch. The continuous probability map quantifies the quality of
the predicted image patch. Inferior quality, associated with lower probability
values, is commonly associated with complex structures, e.g., the cortical
ribbon. Superior quality, associated with higher probability values,
corresponds to flat regions with simple structures. In the encoding path of
the discriminator, the input is downsampled three times; in the decoding path,
the feature maps are upsampled three times. For downsampling/upsampling, we
use a $4\times 4\times 4$ convolution/deconvolution layer, followed by IN and
ReLU activation. The numbers of feature channels are $64$, $128$, and $256$ in
the three stages of the discriminator. Deep supervision strategy [24] is
incorporated in the decoding path to strengthen gradient back propagation.

Figure 5: Network architecture of the multi-scale discriminator.

###  2.2 Loss Functions

We incorporated supervised learning with multi-scale information via deep
supervision strategy [24]. The loss function $\mathcal{L}_{\mathrm{MGAN}}$ is
defined as:

$\mathcal{L}_{\mathrm{MGAN}}=\mathcal{L}_{s_{1}}+\mathcal{L}_{s_{2}}+\mathcal{L}_{s_{3}},$
(4)

where $s_{1}$, $s_{2}$, and $s_{3}$ refer to the three scales employed [31,
32]. For each scale, the objective function is composed of three loss
functions to effectively learn the prediction task. The loss functions are
described next.

####  2.2.1 Adversarial Loss

We propose to use the standard adversarial loss function, which aims to match
the distribution of the predicted images with that of the real images. It is
given by

$\displaystyle\mathcal{L}(G_{a},G_{b},D_{a},D_{b})$
$\displaystyle=\mathbb{E}_{I_{t_{a}}}[\log(D_{a}(I_{t_{a}})]$ (5)
$\displaystyle+\mathbb{E}_{I_{t_{b}}}[\log(1-D_{a}(G_{b}(I_{t_{b}})))]$
$\displaystyle+\mathbb{E}_{I_{t_{b}}}[\log(D_{b}(I_{t_{b}})]$
$\displaystyle+\mathbb{E}_{I_{t_{a}}}[\log(1-D_{b}(G_{a}(I_{t_{a}})))],$

where $I_{t_{a}}$ and $I_{t_{b}}$ refer to the images at time-point $t_{a}$
and $t_{b}$, respectively, $G_{a}$ and $G_{b}$ are the mapping functions, and
$D_{a}$ and $D_{b}$ are the discriminators.

####  2.2.2 Paired Loss

The generators $G_{a}$ and $G_{b}$ seek to minimize the difference between
real and predicted images. We propose to enhance the performance of the
generators by defining a paired loss function that constraints the difference
at voxel-, feature-, and frequency-level. Our paired loss function
$\mathcal{L}_{\mathrm{gen}}^{(\cdot)}$ consists of three loss terms: (i)
quality-driven loss, (ii) texture loss, and (iii) frequency loss.

$\displaystyle\mathcal{L}_{\mathrm{gen}}^{G_{a}}$
$\displaystyle=\mathcal{L}_{\mathrm{Q}}^{G_{a}}+\mathcal{L}_{\mathrm{T}}^{G_{a}}+\mathcal{L}_{\mathrm{F}}^{G_{a}},$
(6) $\displaystyle\mathcal{L}_{\mathrm{gen}}^{G_{b}}$
$\displaystyle=\mathcal{L}_{\mathrm{Q}}^{G_{b}}+\mathcal{L}_{\mathrm{T}}^{G_{b}}+\mathcal{L}_{\mathrm{F}}^{G_{b}}.$

##### Quality-driven loss

The low tissue contrast and the dramatic brain growth hinder translation of
regions such as the convoluted cerebral cortex. Here, we present a quality-
guided learning strategy to strengthen the transformation of the unfathomable
regions. The discriminator outputs a quality map that defines the voxel-wise
probabilities for each predicted image. The heterogeneous distribution of the
probabilities in the quality map motivates us to treat voxels differently.
Voxels with lower probabilities correspond to poor prediction and require more
attention compared to those with higher probabilities. This enhances the image
translation power of the generator at complex regions in the infant brain MRI.
The quality-driven loss $\mathcal{L}_{\mathrm{Q}}^{(\cdot)}$ is defined as:

$\displaystyle\mathcal{L}_{\mathrm{Q}}^{G_{a}}(G_{a};\theta^{G_{a}})=\mathbb{E}_{I_{t_{a}},I_{t_{b}},Q^{D_{b}}}[\|I_{t_{b}}-G_{a}(I_{t_{a}})\|_{1}\odot(1-Q^{D_{b}})^{\beta}],$
(7)
$\displaystyle\mathcal{L}_{\mathrm{Q}}^{G_{b}}(G_{b};\theta^{G_{b}})=\mathbb{E}_{I_{t_{a}},I_{t_{b}},Q^{D_{a}}}[\|I_{t_{a}}-G_{b}(I_{t_{b}})\|_{1}\odot(1-Q^{D_{a}})^{\beta}],$

where $Q^{(\cdot)}$ is the quality map, $\theta^{(\cdot)}$ denotes the
parameters of the network, $\odot$ defines the element-wise multiplication and
$\beta$ represents the parameter that enables to focus on difficult-to-predict
regions. If $\beta$ is set to zero, then $\mathcal{L}_{\mathrm{Q}}^{(\cdot)}$
will be equivalent to $\mathcal{L}_{1}$ norm; losing the ability to define
adaptive weights based on quality map. In this study, we empirically set its
value to $1.5$.

##### Texture loss

This loss ensures that the predicted image has a texture similar to the target
image, and it is defined as the mean square error (MSE) between the Gram
matrix of the target and the predicted image [33, 34]:

$\displaystyle\mathcal{L}_{\mathrm{T}}^{G_{a}}$ $\displaystyle=\Arrowvert
M(I_{t_{b}})-M(G_{a}(I_{t_{a}}))\|_{2},$ (8)
$\displaystyle\mathcal{L}_{\mathrm{T}}^{G_{b}}$ $\displaystyle=\Arrowvert
M(I_{t_{a}})-M(G_{b}(I_{t_{b}}))\|_{2}.$

The gram matrix $M(\cdot)$ is the inner product of the generated images.

##### Frequency loss

The frequency loss $\mathcal{L}_{\mathrm{F}}^{(\cdot)}$ is incorporated via
wavelet decomposition of the generators’ outputs, which steers the effective
prediction of the structural details. $\mathcal{L}_{\mathrm{F}}^{(\cdot)}$ is
defined as:

$\displaystyle\mathcal{L}_{\mathrm{F}}^{G_{a}}$ $\displaystyle=\sum_{k\in
K}\Arrowvert\mathrm{DWT}(I_{t_{b}})^{k}-\mathrm{DWT}(G_{a}(I_{t_{a}}))^{k}\lVert_{1},$
(9) $\displaystyle\mathcal{L}_{\mathrm{F}}^{G_{b}}$ $\displaystyle=\sum_{k\in
K}\Arrowvert\mathrm{DWT}(I_{t_{a}})^{k}-\mathrm{DWT}(G_{b}(I_{t_{b}}))^{k}\lVert_{1},$

where $K=\allowbreak\\{LLL,LLH,LHL,HLL,LHH,HLH,HHL,\allowbreak HHH\\}$; $LLL$
corresponds to the approximation coefficients which encode the image contrast,
and the remaining terms correspond to the detail coefficients, encoding the
high frequency structural details. The wavelet coefficients are decomposed
using bior1.3 [35], which is compactly supported by a biorthogonal spline
wavelet [36].

####  2.2.3 Cycle Consistency Loss

The cycle consistency loss function $\mathcal{L}_{\mathrm{cyc}}$ ensures that
the image prediction cycle brings the predicted image back to the original
image, i.e., $G_{b}(G_{a}(I_{t_{a}}))\approx I_{t_{a}}$ and it is given by:

$\displaystyle\mathcal{L}_{\mathrm{cyc}}(G_{a},G_{b})$
$\displaystyle=\mathbb{E}_{I_{t_{a}}}[\Arrowvert
I_{t_{a}}-G_{b}(G_{a}(I_{t_{a}}))\lVert_{1}],$ (10)
$\displaystyle+\mathbb{E}_{I_{t_{b}}}[\Arrowvert
I_{{t_{b}}}-G_{a}(G_{b}(I_{{t_{b}}}))\lVert_{1}].$

This loss function constraints both the forward and backward image prediction
cycles, causing $G_{a}$ and $G_{b}$ to be consistent with each other.

##  3 Experimental Results

###  3.1 Data Acquisition and Preprocessing

The dataset consists of longitudinal T1-weighted (T1w) and T2-weighted (T2w)
MR images of healthy infant subjects enrolled in the Multi-visit Advanced
Pediatric Brain Imaging (MAP) study. Informed written consent was obtained
from the parents of all the participants and all study protocols were approved
by the University of North Carolina at Chapel Hill Institutional Review Board.
Each subject was scanned every three months in the first postnatal year. The
imaging parameters for T1w MRI data were: TR $=1900$ ms, TE $=4.38$ ms, flip
angle $=7^{\circ}$. All the images had $144$ sagittal slices and $1$ mm
isotropic voxel resolution. The imaging parameters for T2w MR images were TR =
7380 ms, TE = 119 ms, flip angle = $150^{\circ}$, $64$ sagittal slices, and
$1.25\times 1.25\times 1.95\,\text{mm}^{3}$ voxel size.

The dataset was preprocessed using our infant-dedicated preprocessing pipeline
[37, 38]. Then, all the postnatal images of each subject were linearly aligned
to their corresponding 12-months-old images and resampled to the size of
$256\times 256\times 256$ with $1\times 1\times 1\,\mathrm{mm}^{3}$ voxel
resolution. We randomly split the MRI data from $30$ healthy infants into 20
and 10 for training and testing, respectively. Five-fold cross-validation was
performed to tune the hyper-parameters.

###  3.2 Implementation Details

The proposed metamorphic GAN was implemented using TensorFlow library [39] on
a single Nvidia TitanX (Pascal) GPU. Adam optimizer [40] was adopted with an
initial learning rate of $1\times 10^{-4}$ and batch size of $1$. Training,
validation, and testing were performed separately for T1w and T2w images.

During training, we uniformly sampled $3$D patches from each image
encompassing the brain region with a dense stride of 10, providing sufficient
samples for training. The generator was first trained with $5$ epochs before
the adversarial training. The adversarial training was stopped at $50$ epochs.

During inference, the $N=20$ inferences were performed for the estimation of
epistemic and aleatoric uncertainty. The keep rate of the dropout layers was
set to $0.8$. Test-time data augmentation was carried out using a combination
of random flip, rotation along each of the three axes, and random noise, which
were modeled respectively with discrete Bernoulli distribution
$\mathcal{B}(0.5)$, uniform distribution $\mathcal{U}(0,2\pi)$, and normal
distribution $\mathcal{N}(0,0.05)$.

###  3.3 Evaluation Criteria

We employed two commonly used metrics to evaluate the quality of the predicted
images: (i) peak signal-to-noise ratio (PSNR), and (ii) structural similarity
(SSIM) [41]. Higher PSNR and SSIM correspond to accurate image prediction.

###  3.4 Comparison with Existing Techniques

We compared MGAN with three widely used GANs: CycleGAN [16], Pix2Pix [7], and
WGAN [42]. All the compared models were used to predict the 12-month-old brain
MRI from the 2-week-old brain MRI. The prediction task is challenging due to
the extent of changes between the two time points (Fig. 1). For fair
comparison, we re-trained the GANs for optimal parameters.

The image prediction results shown for the compared models in Fig. 6(a)
indicate that MGAN yields T1w and T2w image predictions that are closer to the
ground truth with richer details than the other models. The error maps
indicate that MGAN achieves the lowest error among all methods, especially
around the ventricles and cerebral cortex. Summary statistics for PSNR and
SSIM are reported in Table 1. MGAN achieves significant improvement ($p<0.05$,
paired $t$-test) for PSNR and SSIM over other methods.

(a) T1w image predictions Figure 6: Longitudinal image prediction with various
GANs.

(b) T2w image predictions

Table 1: Summary statistics of PSNR and SSIM for different GANs.

| T1w | T2w
---|---|---
Method | PSNR | SSIM (%) | PSNR | SSIM (%)
CycleGAN | 22.6$\pm$1.1 | 74.2$\pm$2.8 | 21.8$\pm$1.0 | 75.3$\pm$2.2
Pix2Pix | 23.0$\pm$1.3 | 76.2$\pm$3.4 | 22.9$\pm$0.9 | 77.2$\pm$2.8
WGAN | 24.1$\pm$1.2 | 79.4$\pm$2.4 | 23.4$\pm$0.9 | 79.5$\pm$2.0
MGAN | 26.4$\pm$0.9 | 84.0$\pm$2.2 | 25.5$\pm$0.7 | 84.8$\pm$1.8

###  3.5 Ablation Study

Here, we investigate the effectiveness of three components of MGAN — the SFT
block, quality-guided learning, and the hybrid loss function. The influence of
frequency transform on longitudinal prediction was verified using two variants
of the metamorphic generator: (i) incorporating the SFT block equipped with
both the frequency and spatial transform branches, and (ii) replacing the SFT
block with a conventional spatial transform branch. We also investigated the
efficacy of quality-guided learning by conducting experiments with/without
quality maps generated by the discriminators. The configurations are
summarized as follows:

  * •

Backbone: SFT with only spatial transform branch and without quality guidance.

  * •

SFT-NCG: SFT without quality guidance.

  * •

ST-CG: Conventional spatial transform and quality guidance.

  * •

MGAN: SFT and quality guidance.

Table 2 indicates that MGAN achieves the highest PSNR and SSIM with a
significant improvement ($p<0.05$, paired $t$-test). SFT-NCG and ST-CG perform
better than Backbone, validating that wavelet-based feature mapping and
quality-guidance improve prediction accuracy.

Table 2: Ablation study with different MGAN configurations.

| T1w | T2w
---|---|---
Model | PSNR | SSIM (%) | PSNR | SSIM (%)
Backbone | 24.9$\pm$0.5 | 80.3$\pm$1.3 | 24.2$\pm$0.4 | 81.1$\pm$1.2
SFT-NCG | 25.2$\pm$1.0 | 83.7$\pm$2.0 | 25.1$\pm$0.9 | 83.0$\pm$1.6
ST-CG | 25.9$\pm$1.2 | 82.5$\pm$2.0 | 24.8$\pm$1.0 | 82.2$\pm$1.5
MGAN | 26.4$\pm$0.9 | 84.0$\pm$2.2 | 25.5$\pm$0.7 | 84.8$\pm$1.8

Fig. 7(a) shows that Backbone and ST-CG predict the 12-month scan poorly due
to the spatial complexity of the cortical ribbon. SFT-NCG generated
unsatisfactory results at difficult-to-predict regions as indicated by the
high values in the error map. MGAN yields the most accurate prediction, which
matches the ground truth both in terms of tissue contrast and anatomical
structure.

(a) T1w image predictions Figure 7: Longitudinal image prediction results
obtained with different MGAN configurations.

(b) T2w image predictions

We investigated the contribution of the uncertainty-aware loss
$\mathcal{L}_{\mathrm{Q}}$, texture loss $\mathcal{L}_{\mathrm{T}}$, and
frequency loss $\mathcal{L}_{\mathrm{F}}$. Table 3 indicates that including
all loss terms (Eq. 6) yields the highest PSNR and SSIM. In contrast, using
only the uncertainty-aware loss yields the lowest PSNR and SSIM. This implies
that both the texture and frequency losses improve the predictive power of the
generator.

Table 3: Ablation study with different combinations of losses.

|  |  | T1w | T2w
---|---|---|---|---
$\mathcal{L}_{\mathrm{Q}}$ | $\mathcal{L}_{\mathrm{T}}$ | $\mathcal{L}_{\mathrm{F}}$ | PSNR | SSIM (%) | PSNR | SSIM (%)
✓ |  |  | 26.0$\pm$1.3 | 83.3$\pm$1.4 | 25.1$\pm$0.8 | 83.1$\pm$1.4
✓ | ✓ |  | 26.2$\pm$1.0 | 83.5$\pm$2.0 | 25.3$\pm$0.9 | 83.7$\pm$1.9
✓ |  | ✓ | 26.1$\pm$1.0 | 83.7$\pm$1.8 | 25.2$\pm$0.8 | 84.3$\pm$1.7
✓ | ✓ | ✓ | 26.4$\pm$0.9 | 84.0$\pm$2.2 | 25.5$\pm$0.7 | 84.8$\pm$1.8

(a) T1w image predictions Figure 8: Longitudinal prediction results for
different time points. (Left) The forward path predicts a 12-month-old image
from images at earlier time points. (Right) The backward path predicts images
of earlier time points from a 12-month-old image.

(b) T2w image predictions

###  3.6 Longitudinal Prediction

We demonstrate the effectiveness MGAN in predicting a 12-month-old image from
any earlier time-point, i.e., 2 weeks, 3 months, 6 months, and 9 months.
Predictions from the forward and backward prediction paths are evaluated. The
predicted images along with the error maps and uncertainty maps are shown in
Fig. 8(a). The quantitative results are presented in Table 4. Despite the
significant differences in appearance and structure, MGAN is able to predict
the images with great resemblance to the ground-truth images in both tissue
contrast and anatomical structure. This is validated by the high PSNR and SSIM
values. The corresponding epistemic and aleatoric uncertainty maps of the
predictions are also shown in Fig. 8(a). The epistemic and aleatoric
uncertainty is positively correlated with prediction errors.

Table 4: Statistical summary of evaluation metrics for longitudinal
prediction.

|  | T1w | T2w
---|---|---|---
|  | PSNR | SSIM(%) | PSNR | SSIM(%)
forward prediction | 0m$\rightarrow$12m | 26.4$\pm$0.9 | 84.0$\pm$2.2 | 25.5$\pm$0.7 | 84.8$\pm$1.8
3m$\rightarrow$12m | 26.1$\pm$1.3 | 84.7$\pm$4.0 | 25.7$\pm$0.9 | 83.8$\pm$2.2
6m$\rightarrow$12m | 27.7$\pm$2.2 | 89.1$\pm$3.6 | 26.8$\pm$1.8 | 87.5$\pm$2.8
9m$\rightarrow$12m | 29.0$\pm$2.9 | 89.9$\pm$4.6 | 28.5$\pm$1.9 | 88.3$\pm$2.8
backward prediction | 12m$\rightarrow$0m | 27.1$\pm$0.9 | 86.7$\pm$0.2 | 26.5$\pm$1.1 | 86.7$\pm$1.2
12m$\rightarrow$3m | 26.9$\pm$1.7 | 86.9$\pm$3.1 | 26.4$\pm$1.2 | 86.4$\pm$2.2
12m$\rightarrow$6m | 27.8$\pm$1.7 | 89.7$\pm$2.7 | 27.1$\pm$1.8 | 89.4$\pm$3.2
12m$\rightarrow$9m | 28.4$\pm$2.5 | 90.5$\pm$3.1 | 27.6$\pm$2.2 | 90.1$\pm$3.4

Figure 9: Quality visualization for the 0-to-12-month-old prediction.

##  4 Discussion

In this paper, we presented a metamorphic GAN that can be trained to predict
infant brain MRI from one time point to another. Longitudinal prediction of
infant brain MRI is challenging owing to rapid contrast and structural changes
in the first year of life. To capture these changes, our MGAN incorporates a
SFT block and integrates quality-guided learning via a hybrid loss function.

We compared our method with existing generative adversarial networks, such as
CycleGAN, Pix2Pix, and WGAN. We found that these networks are effective in
prediction structures at a global scale but are less effective in predicting
fine-scale structural details, especially in the cortex (Fig. 6(a)). In
contrast, our prediction network capture spatially heterogeneous changes by
employing both spatial and frequency transforms to generate feature maps.
Particularly, DWT-based frequency transform decomposes the image into low and
high frequency components to help the translation of image contrast and subtle
details (Fig. 7(a)).

The quality-guided learning strategy involves using an estimation map for
characterizing voxel-wise prediction quality. Fig. 9 shows that regions with
complex structures, e.g., the cerebral cortex, are associated with higher bias
values. In contrast, regions with simple structure, e.g., lateral ventricles,
are associated with lower bias values. As shown in Fig. 7(a), employing the
quality-driven $\mathcal{L}_{Q}$ loss results in more accurate predictions at
challenging regions with complex patterns. Additionally, the wavelet
decomposition and gram matrix enhance the similarity between predictions and
ground truths both in terms of content and style (Table 3).

##  5 Conclusion

We have proposed a trustworthy learning-based framework for longitudinal
postnatal brain MRI prediction. The key feature our method is the utilization
of wavelet transform to enable image prediction at multiple frequencies. We
utilize quality guidance to strengthen the learning of prediction of
challenging regions. We employ a hybrid loss function and a multi-scale
discriminator to capture differences in global intensity, style, and
structure. Experimental results demonstrate that our method achieves superior
performance over several state-of-the-art image-to-image translation networks.
Despite the effectiveness of our method, it is currently trained with paired
data. In future, it can be extended to be trainable with unpaired data.

## Acknowledgments

This work was supported in part by United States National Institutes of Health
(NIH) grants EB008374, EB006733, and AG053867. Y. Huang was supported by the
China Scholarship Council and the National Natural Science Foundation of China
under Grant 6210011424.

## References

  * [1] J. Dubois, M. Benders, A. Cachia, F. Lazeyras, R. Ha-Vinh Leuchter, S. V. Sizonenko, C. Borradori-Tolsa, J.-F. Mangin, P. S. Hüppi, Mapping the early cortical folding process in the preterm newborn brain, Cerebral Cortex 18 (6) (2008) 1444 – 1454.
  * [2] B. R. Howell, M. A. Styner, W. Gao, P.-T. Yap, L. Wang, K. Baluyot, E. Yacoub, G. Chen, T. Potts, A. Salzwedel, G. Li, J. H. Gilmore, J. Piven, J. K. Smith, D. Shen, K. Ugurbil, H. Zhu, W. Lin, J. T. Elison, The UNC/UMN baby connectome project (BCP): An overview of the study design and protocol development, NeuroImage 185 (2019) 891 – 905.  doi:https://doi.org/10.1016/j.neuroimage.2018.03.049.
URL http://www.sciencedirect.com/science/article/pii/S1053811918302593

  * [3] J. H. Gilmore, R. C. Knickmeyer, W. Gao, Imaging structural and functional brain development in early childhood, Nature Reviews Neuroscience 19 (3) (2018) 123.
  * [4] J. Matsuzawa, M. Matsui, T. Konishi, K. Noguchi, R. C. Gur, W. Bilker, T. Miyawaki, Age-related volumetric changes of brain gray and white matter in healthy infants and children, Cerebral Cortex 11 (4) (2001) 335 – 342.  doi:https://doi.org/10.1093/cercor/11.4.335.
  * [5] R. C. Knickmeyer, S. Gouttard, C. Kang, D. Evans, K. Wilber, J. K. Smith, R. M. Hamer, W. Lin, G. Gerig, J. H. Gilmore, A structural MRI study of human brain development from birth to 2 years, Journal of Neuroscience 28 (47) (2008) 12176 – 12182.
  * [6] T. Paus, D. Collins, A. Evans, G. Leonard, B. Pike, A. Zijdenbos, Maturation of white matter in the human brain: a review of magnetic resonance studies, Brain research bulletin 54 (3) (2001) 255 – 266.
  * [7] P. Isola, J.-Y. Zhu, T. Zhou, A. A. Efros, Image-to-image translation with conditional adversarial networks, 2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR) (2016) 5967 – 5976.
  * [8] J.-Y. Zhu, T. Park, P. Isola, A. A. Efros, Unpaired image-to-image translation using cycle-consistent adversarial networks, 2017 IEEE International Conference on Computer Vision (ICCV) (2017) 2242 – 2251.
  * [9] X. Huang, M.-Y. Liu, S. J. Belongie, J. Kautz, Multimodal unsupervised image-to-image translation, in: ECCV, 2018.
  * [10] M.-Y. Liu, T. Breuel, J. Kautz, Unsupervised image-to-image translation networks, in: NIPS, 2017.
  * [11] I. J. Goodfellow, J. Pouget-Abadie, M. Mirza, B. Xu, D. Warde-Farley, S. Ozair, A. C. Courville, Y. Bengio, Generative adversarial networks, ArXiv abs/1406.2661.
  * [12] D. Nie, R. Trullo, J. Lian, C. Petitjean, S. Ruan, Q. Wang, D. Shen, Medical image synthesis with context-aware generative adversarial networks, in: International Conference on Medical Image Computing and Computer-Assisted Intervention, Springer, 2017, pp. 417 – 425.
  * [13] D. Nie, D. Shen, Adversarial confidence learning for medical image segmentation and synthesis, International Journal of Computer Vision (2020) 1 – 20.
  * [14] K. Lee, M.-K. Choi, H. Jung, DavinciGAN: Unpaired surgical instrument translation for data augmentation, in: MIDL, 2018.
  * [15] K. Armanious, C. Yang, M. Fischer, T. Küstner, K. Nikolaou, S. Gatidis, B. Yang, MedGAN: Medical image translation using GANs, Computerized medical imaging and graphics : the official journal of the Computerized Medical Imaging Society 79 (2019) 101684.
  * [16] Z. Zhang, L. Yang, Y. Zheng, Translating and segmenting multimodal medical volumes with cycle- and shape-consistency generative adversarial network, 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition (2018) 9242 – 9251.
  * [17] A. Chartsias, T. Joyce, R. Dharmakumar, S. A. Tsaftaris, Adversarial image synthesis for unpaired multi-modal cardiac data, in: SASHIMI@MICCAI, 2017.
  * [18] D. Nie, R. Trullo, J. Lian, L. Wang, C. Petitjean, S. Ruan, Q. Wang, D. Shen, Medical image synthesis with deep convolutional adversarial networks, IEEE Transactions on Biomedical Engineering 65 (12) (2018) 2720 – 2730.
  * [19] Y. Hiasa, Y. Otake, M. Takao, T. Matsuoka, K. Takashima, A. Carass, J. L. Prince, N. Sugano, Y. Sato, Cross-modality image synthesis from unpaired data using CycleGAN, in: International workshop on simulation and synthesis in medical imaging, Springer, 2018, pp. 31 – 41.
  * [20] H. Yang, J. Sun, A. Carass, C. Zhao, J. Lee, Z. Xu, J. L. Prince, Unpaired brain MR-to-CT synthesis using a structure-constrained CycleGAN, in: DLMIA/ML-CDS@MICCAI, 2018.
  * [21] D. Ulyanov, A. Vedaldi, V. S. Lempitsky, Instance normalization: The missing ingredient for fast stylization, ArXiv abs/1607.08022.
  * [22] X. Glorot, A. Bordes, Y. Bengio, Deep sparse rectifier neural networks, in: AISTATS, 2011.
  * [23] K. He, X. Zhang, S. Ren, J. Sun, Deep residual learning for image recognition, 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR) (2016) 770 – 778.
  * [24] S. Xie, Z. Tu, Holistically-nested edge detection, International Journal of Computer Vision 125 (2015) 3–18.
  * [25] A. Der Kiureghian, O. Ditlevsen, Aleatory or epistemic? does it matter?, Structural safety 31 (2) (2009) 105–112.
  * [26] M. Abdar, F. Pourpanah, S. Hussain, D. Rezazadegan, L. Liu, M. Ghavamzadeh, P. Fieguth, A. Khosravi, U. R. Acharya, V. Makarenkov, et al., A review of uncertainty quantification in deep learning: Techniques, applications and challenges, arXiv preprint arXiv:2011.06225.
  * [27] A. Kendall, V. Badrinarayanan, R. Cipolla, Bayesian segnet: Model uncertainty in deep convolutional encoder-decoder architectures for scene understanding, arXiv preprint arXiv:1511.02680.
  * [28] Y. Gal, Z. Ghahramani, Dropout as a bayesian approximation: Representing model uncertainty in deep learning, in: international conference on machine learning, PMLR, 2016, pp. 1050–1059.
  * [29] M. S. Ayhan, P. Berens, Test-time data augmentation for estimation of heteroscedastic aleatoric uncertainty in deep neural networks.
  * [30] G. Wang, W. Li, M. Aertsen, J. Deprest, S. Ourselin, T. Vercauteren, Aleatoric uncertainty estimation with test-time augmentation for medical image segmentation with convolutional neural networks, Neurocomputing 338 (2019) 34–45.
  * [31] T.-C. Wang, M.-Y. Liu, J.-Y. Zhu, A. Tao, J. Kautz, B. Catanzaro, High-resolution image synthesis and semantic manipulation with conditional GANs, in: Proceedings of the IEEE conference on computer vision and pattern recognition, 2018, pp. 8798 – 8807.
  * [32] A. Karnewar, O. Wang, MSG-GAN: Multi-scale gradients for generative adversarial networks, in: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2020, pp. 7799 – 7808.
  * [33] L. Gatys, A. S. Ecker, M. Bethge, Texture synthesis using convolutional neural networks, in: Advances in neural information processing systems, 2015, pp. 262 – 270.
  * [34] Y. Li, C. Fang, J. Yang, Z. Wang, X. Lu, M.-H. Yang, Diversified texture synthesis with feed-forward networks, in: Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2017, pp. 3920 – 3928.
  * [35] A. Cohen, Biorthogonal bases of compactly supported wavelets, 2006.
  * [36] R. Szewczyk, K. Grabowski, M. Napieralska, W. Sankowski, M. Zubert, A. Napieralski, A reliable iris recognition algorithm based on reverse biorthogonal wavelet transform, Pattern Recognition Letters 33 (8) (2012) 1019 – 1026.
  * [37] Y. Dai, F. Shi, L. Wang, G. Wu, D. Shen, iBEAT: a toolbox for infant brain magnetic resonance image processing, Neuroinformatics 11 (2) (2013) 211 – 225\.
  * [38] G. Li, J. Nie, L. Wang, F. Shi, W. Lin, J. H. Gilmore, D. Shen, Mapping region-specific longitudinal cortical surface expansion from birth to 2 years of age, Cerebral cortex 23 (11) (2013) 2724 – 2733.
  * [39] M. Abadi, P. Barham, J. Chen, Z. Chen, A. Davis, J. Dean, M. Devin, S. Ghemawat, G. Irving, M. Isard, M. Kudlur, J. Levenberg, R. Monga, S. Moore, D. Murray, B. Steiner, P. Tucker, V. Vasudevan, P. Warden, X. Zhang, TensorFlow: A system for large-scale machine learning.
  * [40] D. P. Kingma, J. Ba, Adam: A method for stochastic optimization, CoRR abs/1412.6980.
  * [41] Z. Wang, A. C. Bovik, H. R. Sheikh, E. P. Simoncelli, Image quality assessment: from error visibility to structural similarity, IEEE Transactions on Image Processing 13 (2004) 600 – 612.
  * [42] M. Arjovsky, S. Chintala, L. Bottou, Wasserstein generative adversarial networks, in: ICML, 2017.