# Retinal Image Restoration using Transformer and Cycle-Consistent Generative
Adversarial Network

Alnur Alimanov  Department of Computer Engineering
Bahcesehir University, Istanbul, Turkey
Email<EMAIL_ADDRESS>Md Baharul Islam  Bahcesehir
University, Istanbul, Turkey
American University of Malta
ORCID: 0000-0002-9928-5776

###### Abstract

Medical imaging plays a significant role in detecting and treating various
diseases. However, these images often happen to be of too poor quality,
leading to decreased efficiency, extra expenses, and even incorrect diagnoses.
Therefore, we propose a retinal image enhancement method using a vision
transformer and convolutional neural network. It builds a cycle-consistent
generative adversarial network that relies on unpaired datasets. It consists
of two generators that translate images from one domain to another (e.g., low-
to high-quality and vice versa), playing an adversarial game with two
discriminators. Generators produce indistinguishable images for discriminators
that predict the original images from generated ones. Generators are a
combination of vision transformer (ViT) encoder and convolutional neural
network (CNN) decoder. Discriminators include traditional CNN encoders. The
resulting improved images have been tested quantitatively using such
evaluation metrics as peak signal-to-noise ratio (PSNR), structural similarity
index measure (SSIM), and qualitatively, i.e., vessel segmentation. The
proposed method successfully reduces the adverse effects of blurring, noise,
illumination disturbances, and color distortions while significantly
preserving structural and color information. Experimental results show the
superiority of the proposed method. Our testing PSNR is 31.138 dB for the
first and 27.798 dB for the second dataset. Testing SSIM is 0.919 and 0.904,
respectively. The code is available at https://github.com/AAleka/Transformer-
Cycle-GAN

###### Index Terms:

Retinal image restoration, deep learning, transformer, generative adversarial
network.

##  I Introduction

The retina of the central nervous system is responsible for transforming the
incoming light into a neural signal and sending it to the brain’s visual
cortex for processing. Therefore, ophthalmologists can diagnose eye, brain,
and blood circulation diseases by analyzing the retinal images. However,
sometimes due to poor imaging device configuration, misoperations during the
imaging process, and patient restlessness (primarily applicable to children),
these images suffer from such degradation types as out-of-focus blurring,
image noise, low, high, and uneven illuminations, and color distortion. It
suffers a tremendous negative impact on treatment efficiency, the extra time
and money expenditure, and even misdiagnosis.

In recent years, deep learning-based techniques [1, 2, 3] are used to detect
and classify diseases using retinal images. However, these methods rely on
high-quality images during analysis, which is not always possible to acquire
in some cases. Researchers also propose traditional methods [4, 5] to enhance
these images. However, they often suffer from generalization issues, meaning
they cannot be applied to all cases. For example, Zhang et al. [5], took the
reflective feature of the fundus camera into account. They proposed a double-
pass fundus reflection model (DPFR) that improves the contrast of retinal
images. However, the presence of image artifacts is noticeable. Many learning-
based techniques [6, 7, 8, 9, 10, 11] have been introduced to solve this
problem. However, most of these works solved the missing paired data problem
by manually degrading original high-quality fundus images. Manually degraded
images are different from original low-quality retinal images. Thus, these
methods show poor performance when enhancing real degradation effects. For
instance, the authors of [11] developed a restoration deep learning network
that was trained using simulated cataract-like images. However, it performs
worse when it comes to restoring real low-quality images.

Figure 1: The workflow of the proposed method. Low-quality images go through
Generator H and the result - Generated HQ - is inputted into Generator H and
Discriminator H. Next, the Generated LQ image is made from high-quality image
using Generator L, then this generated image is passed to both Generator H and
Discriminator L. After these steps, losses are calculated.

This paper introduces a novel deep learning retinal image restoration model
based on a cycle-consistent generative adversarial network with modified
generators that do not require paired datasets. Our model replaced the
traditional CNN encoder with a vision transformer encoder, resulting in faster
convergence, superior quantitative and qualitative outcomes, and better
structural and color preservation than other methods. The method has been
trained and validated using two publicly available datasets [12, 13]. The
contributions can be summarized as follows:

  * •

We propose a novel CycleGAN architecture consisting of a vision transformer
and convolutional neural network. To the best of our knowledge, this is the
first attempt to combine CNN with a transformer in CycleGAN architecture.

  * •

Replacing the transformer decoder with a CNN decoder reduced the computational
expense compared to pure transformer CycleGAN architecture, which allowed us
to use a large image size.

##  II Methodology

###  II-A Model Architecture

Our method is based on CycleGAN, initially proposed by [14]. However, we
replaced the CNN encoder in the generator network with a vision transformer
encoder, as shown in Fig. 1. Using CNN as a decoder decreased computational
expenses compared to a pure transformer network which allowed to increase the
image size by a factor of 4, from 64$\times$64 to 256$\times$256\. This
network consists of a pair of generator networks ($G_{H}$ and $G_{L}$) and a
pair of discriminators ($D_{H}$ and $D_{L}$). The task of each generator is to
translate images from one domain to another and to fool the corresponding
discriminator. For instance, $G_{H}$ is trying to generate realistic, high-
quality images from original low-quality ones, and $D_{H}$ receives this
generated image and one original high-quality image and calculates the
probability of a given image being authentic. As a result, generators and
discriminators are playing an adversarial game.

A discriminator is a CNN-based classifier that consists of such layers as 2D
convolution, LeakyReLU activation function, and instance normalization.
Firstly, a colorful 256$\times$256 image is fed into the convolutional layer
with output channels, kernels size, stride, and padding equal to 64, 4, 2, and
1, respectively. Then, we used the LeakyReLU activation function. The result
sends to three consecutive downsampling blocks: convolutional layer, instance
normalization, and LeakyReLU. Kernel size and padding are equal to 4 and 1,
but output features double after each downsampling process starting from 64
and ending with 512. Stride is 1 in the first and 2 in the last two blocks. In
the end, it is fed into the final convolutional layer with output channel,
stride, and padding equal to 1 and kernel size being 4. As a result, we get a
30$\times$30 patch and make predictions using the sigmoid activation function.

The generator consists of a vision transformer encoder and a CNN decoder.
Firstly, each input image of size 256$\times$256 is divided into patches of
size 8$\times$8 using a convolutional projection proposed by [15]. It provides
additional efficiency. Instead of a depth-wise convolutional layer with batch
normalization, we only used a convolutional layer with 1024 output channels,
kernel, and stride sizes of 1. Then we flattened the 3D output to 2D and
transposed it. It reduced the computation time without affecting the accuracy
of the method. Then these flattened sequences are fed into a standard
transformer encoder network as shown in Fig. 1. Our network’s depth and
projection dimensions are 7 and 1024, respectively. The output shape of this
encoder is 1024$\times$1024\. Therefore, we reshaped it into
1024$\times$32$\times$32 to feed into the CNN decoder. It consists of 3
upsampling blocks: transpose convolution, instance normalization, and ReLU
activation function. The output channels have been decreased by 2 after each
block; kernel size, stride, padding, and output padding in all blocks are 3,
2, 1, and 1, respectively. The resulted matrix of shape
128$\times$256$\times$256 is converted into a color image using the last
convolutional layer with a kernel size of 7, the stride of 1, and padding of
3. The final output is activated using a hyperbolic tangent.

###  II-B Loss functions

####  II-B1 Adversarial Loss

Adversarial loss has been developed in the first GAN work [16] that can be
formulated as discriminator loss. Given low-quality image $L$, high-quality
image $H$, a generator $G_{H}$ that transforms $L$ to high quality $\hat{H}$,
$G_{L}$ that degrades images (from $Y$ to $\hat{X}$), $D_{H}$ and $D_{L}$
classify high- ($H$, $\hat{H}$) and low-quality images ($L$ and $\hat{L}$),
respectively. Mathematically, this loss function is expressed as follows:

$\displaystyle L_{A}(D_{L},L,\hat{L})$
$\displaystyle=E_{L}[log(D_{L}(L))]+E_{\hat{L}}[log(1-D_{L}(\hat{L}))]$
$\displaystyle L_{A}(D_{H},H,\hat{H})$
$\displaystyle=E_{H}[log(D_{H}(H))]+E_{\hat{H}}[log(1-D_{H}(\hat{H}))]$ (1)

####  II-B2 Cycle Consistency Loss

Cycle-consistency can be formulated as follows: a generated image $\hat{H}$
translated back to $\acute{L}$ should not be different from $L$. The same
things apply for $H$ and a regenerated $\acute{H}$ from $\hat{L}$. The
difference between regenerated and original images is calculated using the
$L1$ function in the following way:

$\displaystyle L_{C}(\acute{L},L)$ $\displaystyle=E_{L}[||\acute{L}-L||_{1}]$
$\displaystyle L_{C}(\acute{H},H)$ $\displaystyle=E_{H}[||\acute{H}-H||_{1}]$
(2)

####  II-B3 Identity loss

To calculate identity loss, we use generators and images from the same
domains. More specifically, $H$ should not be modified when it is sent to
$G_{H}$. The same things apply to $L$ and $G_{L}$. To calculate this loss
function, we also use the $L1$ function in the following way:

$\displaystyle L_{I}(L,G_{L}(L))$ $\displaystyle=E_{L}[||L-G_{L}(L)||_{1}]$
$\displaystyle L_{I}(H,G_{H}(H))$ $\displaystyle=E_{H}[||H-G_{H}(H)||_{1}]$
(3)

####  II-B4 Total loss

The total generator and discriminator loss can be formulated as:

$\displaystyle L(G_{L},G_{H},D_{L},D_{H})$
$\displaystyle=L_{A}(D_{L},L,\hat{L})+L_{A}(D_{H},H,\hat{H})+$
$\displaystyle+\lambda_{1}L_{C}(\acute{L},\hat{L})+\lambda_{1}L_{C}(\acute{H},\hat{H})+$
$\displaystyle+\lambda_{2}L_{I}(L,G_{L}(L))+\lambda_{2}L_{I}(H,G_{H}(H))$ (4)

where $\lambda_{1}$ and $\lambda_{2}$ are weights of cycle consistency and
identity losses. The main problem that discriminators and generators are
trying to solve can be mathematically expressed in the following equation:

$G_{L}^{*},G_{H}^{*}=arg\min_{G_{L},G_{H}}\max_{D_{L},D_{H}}L(G_{L},G_{H},D_{L},D_{H})$
(5)

##  III Results and Discussion

###  III-A Datasets

In our work, we used two datasets, such as EyeQ [12] and Mendeley [13]. EyeQ
dataset has 16818 ”Good”, 6434 ”Usable”, and 5,538 ”Bad” images of size
800$\times$800\. The Mendeley dataset comprises 1146 ”Artifacts” and 1060 ”No
artifacts” images of size 350$\times$346\. For training, we used all ”Good,”
”No artifacts,” 5455 ”Usable,” and 986 ”Artifacts” images. The testing set
includes 979 ”Usable” and 160 ”Artifacts” images. To train vessel segmentation
model, we used DRIVE [17] dataset that includes 40 images of size
565$\times$584 with corresponding manually annotated segmentation’s.

###  III-B Experimental Setup

We implement the proposed model in the PyTorch framework. The hardware
configurations are Intel Core i7-10700f CPU, 32 GB of RAM, and an NVIDIA
GeForce RTX 2080 SUPER 8 GB. The model was set to train for 100 epochs with
early stopping as soon as average generator loss stops decreasing. The
generator loss has decreased, while discriminator loss has increased for 23
epochs. The generator model improved faster than the discriminator until both
converged. The total training time is 25 hours for our experiment.

Figure 2: Comparison with state-of-the-art methods. From top to bottom are
fundus image, corresponding segmentation and zoomed segmentation. (a) low-
quality image, (b) CLAHE [4], (c) CycleGAN [6], (d) CycleCBAM [10], (e) ArcNet
[11], (f) DPFR [5], (g) Ours.

###  III-C Qualitative and Quantitative Performance

To validate the efficiency of our method, we performed qualitative analysis
for restored images using vessel segmentation with 5 state-of-the-art methods:
CLAHE [4], CycleGAN [6], Cycle-CBAM [10], ArcNet [11], DPFR [5]. UNet [18] was
trained using pairs of fundus images with corresponding manually segmented
images from DRIVE dataset [17]. Fig. 2 shows the restoration and vessel
segmentation results. Our method enhanced the retinal image better compared to
the state-of-the-art. In addition, our result has less noise compared to
CycleGAN [6] and CycleCBAM [10]. Algorithm-based methods [4, 5] could improve
the contrast making vessels visible, but produced images have artifacts.
ArcNet [11] failed to restore tiny blood vessels in the dark region.

To further evaluate our results, we conducted the quantitative analysis with
the same methods. The evaluation metrics include PSNR, SSIM, and single image
test time (SITT), as shown in TABLE I. As we can see, our method significantly
outperforms all other methods in terms of PSNR and SSIM with the EyeQ dataset.
Our approach also shows the best results compared to the state-of-the-art for
the Mendeley dataset. In terms of SITT, it shows competitive outcomes taking
the third position after CLAHE [4] and original CycleGAN [6].

TABLE I: Comparison with state-of-the-art methods. Dataset | Method | PSNR | SSIM | SITT
---|---|---|---|---
EyeQ [12] | CLAHE [4] | 25.152 dB | 0.716 | 17ms
CycleGAN [6] | 25.577 dB | 0.882 | 90ms
Cycle-CBAM [10] | 26.263 dB | 0.901 | 100ms
ArcNet [11] | 21.073 dB | 0.87 | 183ms
DPFR [5] | 12.423 dB | 0.34 | 449ms
Ours | 31.138 dB | 0.919 | 97ms
Mendeley [13] | CLAHE [4] | 26.07 dB | 0.77 | 17ms
CycleGAN [6] | 26.08 dB | 0.9 | 90ms
Cycle-CBAM [10] | 25.808 dB | 0.901 | 100ms
ArcNet [11] | 22.17 dB | 0.891 | 183ms
DPFR [5] | 12.797 dB | 0.304 | 449ms
Ours | 27.798 dB | 0.904 | 97ms

###  III-D Ablation study

To conduct an ablation study of our work, we trained and tested the original
fully convolutional CycleGAN model with the same datasets. Total training time
for the traditional CycleGAN required 40 hours to converge, 15 hours more than
ours. Additionally, testing PSNR values of our and traditional methods are
31.138 dB and 25.577 dB for the EyeQ dataset, and for the Mendeley dataset,
the values are 27.798 dB and 26.08 dB, respectively. The SSIM values are 0.919
and 0.882 for the EyeQ and 0.904 and 0.9 for the Mendeley datasets,
respectively. On average, testing one image takes 90ms for CNN CycleGAN and
97ms for our method, meaning that our modification did not add much
computational expense. Fig. 3 demonstrates the qualitative comparison of these
two methods. Our method is better at restoring tiny, barely visible blood
vessels and preserving structural and color information of original images.

Figure 3: Qualitative comparison of tiny blood vessels restoration.

##  IV Conclusion

In this paper, we combined a vision transformer and convolutional neural
network, resulting in better retinal image enhancement than fully
convolutional CycleGAN and other state-of-the-art techniques. In addition, our
goal was to reduce the computational expense of transformer CycleGAN
architecture without compromising the performance. This simple yet effective
combination has led to comparatively better quantitatively, qualitatively, and
computational performance. The model was tested using two datasets, such as
EyeQ [12] and Mendeley dataset [13]. Testing PSNR is 31.138 dB for the first
and 27.798 dB for the second dataset. Testing SSIM is 0.919 and 0.904,
respectively. The ablation study demonstrates a significant advantage of our
method compared to the original CycleGAN architecture.

## References

  * [1] F. Grassmann, J. Mengelkamp, C. Brandl, S. Harsch, M. E. Zimmermann, B. Linkohr, A. Peters, I. M. Heid, C. Palm, and B. H. Weber, “A deep learning algorithm for prediction of age-related eye disease study severity scale for age-related macular degeneration from color fundus photography,” _Ophthalmology_ , vol. 125, no. 9, pp. 1410–1420, 2018.
  * [2] Z. Gao, J. Li, J. Guo, Y. Chen, Z. Yi, and J. Zhong, “Diagnosis of diabetic retinopathy using deep neural networks,” _IEEE Access_ , vol. 7, pp. 3360–3370, 2018.
  * [3] Y. Peng, S. Dharssi, Q. Chen, T. D. Keenan, E. Agrón, W. T. Wong, E. Y. Chew, and Z. Lu, “Deepseenet: a deep learning model for automated classification of patient-based age-related macular degeneration severity from color fundus photographs,” _Ophthalmology_ , vol. 126, no. 4, pp. 565–575, 2019.
  * [4] F. M. Shamsudeen and G. Raju, “Enhancement of fundus imagery,” in _2016 International Conference on Next Generation Intelligent Systems (ICNGIS)_. IEEE, 2016, pp. 1–5.
  * [5] S. Zhang, C. A. Webers, and T. T. Berendschot, “A double-pass fundus reflection model for efficient single retinal image enhancement,” _Signal Processing_ , vol. 192, p. 108400, 2022.
  * [6] Q. You, C. Wan, J. Sun, J. Shen, H. Ye, and Q. Yu, “Fundus image enhancement method based on cyclegan,” in _2019 41st annual international conference of the IEEE engineering in medicine and biology society (EMBC)_. IEEE, 2019, pp. 4500–4503.
  * [7] W. Li, G. Liu, Y. He, J. Wang, W. Kong, and G. Shi, “Quality improvement of adaptive optics retinal images using conditional adversarial networks,” _Biomedical optics express_ , vol. 11, no. 2, pp. 831–849, 2020.
  * [8] A. D. Pérez, O. Perdomo, H. Rios, F. Rodríguez, and F. A. González, “A conditional generative adversarial network-based method for eye fundus image quality enhancement,” in _International Workshop on Ophthalmic Medical Image Analysis_. Springer, 2020, pp. 185–194.
  * [9] A. Alimanov and M. B. Islam, “Retinal image restoration and vessel segmentation using modified cycle-cbam and cbam-unet,” 2022.
  * [10] C. Wan, X. Zhou, Q. You, J. Sun, J. Shen, S. Zhu, Q. Jiang, and W. Yang, “Retinal image enhancement using cycle-constraint adversarial network,” _Frontiers in Medicine_ , vol. 8, 2021.
  * [11] H. Li, H. Liu, Y. Hu, R. Higashita, Y. Zhao, H. Qi, and J. Liu, “Restoration of cataract fundus images via unsupervised domain adaptation,” in _2021 IEEE 18th International Symposium on Biomedical Imaging (ISBI)_. IEEE, 2021, pp. 516–520.
  * [12] H. Fu, B. Wang, J. Shen, S. Cui, Y. Xu, J. Liu, and L. Shao, “Evaluation of retinal image quality assessment networks in different color-spaces,” in _International Conference on Medical Image Computing and Computer-Assisted Intervention_. Springer, 2019, pp. 48–56.
  * [13] T. K. Yoo, J. Y. Choi, and H. K. Kim, “Cyclegan-based deep learning technique for artifact reduction in fundus photography,” _Graefe’s Archive for Clinical and Experimental Ophthalmology_ , vol. 258, no. 8, pp. 1631–1637, 2020\.
  * [14] J.-Y. Zhu, T. Park, P. Isola, and A. A. Efros, “Unpaired image-to-image translation using cycle-consistent adversarial networks,” in _Proceedings of the IEEE international conference on computer vision_ , 2017, pp. 2223–2232.
  * [15] H. Wu, B. Xiao, N. Codella, M. Liu, X. Dai, L. Yuan, and L. Zhang, “Cvt: Introducing convolutions to vision transformers,” in _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , 2021, pp. 22–31.
  * [16] I. Goodfellow, J. Pouget-Abadie, M. Mirza, B. Xu, D. Warde-Farley, S. Ozair, A. Courville, and Y. Bengio, “Generative adversarial nets,” _Advances in neural information processing systems_ , vol. 27, 2014.
  * [17] J. Staal, M. D. Abràmoff, M. Niemeijer, M. A. Viergever, and B. Van Ginneken, “Ridge-based vessel segmentation in color images of the retina,” _IEEE transactions on medical imaging_ , vol. 23, no. 4, pp. 501–509, 2004.
  * [18] O. Ronneberger, P. Fischer, and T. Brox, “U-net: Convolutional networks for biomedical image segmentation,” in _International Conference on Medical image computing and computer-assisted intervention_. Springer, 2015, pp. 234–241.