Super-resolution Reconstruction Based on Capsule Generative Adversarial Network

The generator network contains two parts: a dense residual network layer and a sub-pixel convolutional layer. The dense residual network contains 16 residual blocks composed of a convolutional layer and an activation function—each residual dense block structure, As shown in Fig. 2. using this structure can avoid the disappearance of the gradient during backpropagation while using Global Residual learning to obtain global dense features from the original LR image, and further combine the shallow features and deep features to fully invoke the layering of low-resolution images features to improve feature extraction capabilities so that the network learns more effective features.

Figure 2.

The discriminator network uses a dual-route capsule network to achieve true or false classifications, including a capsule vector forming part and a dual-routing part. The capsule vector forming part includes a Conv layer, a PrimaryCaps layer, and two DigitCaps layers.

The capsule vector formation part is that after the Conv layer performs preliminary feature extraction on the input image, the original scalar neuron of the convolutional network is repackaged into a vector neuron during the parameter transfer process PrimaryCaps layer. The length of each vector represents the estimated probability of whether the object exists, and its direction records the posture parameters of the object.

After the PrimaryCaps layer, the obtained multi-dimensional tensor is flattened into a one-dimensional array and then input into the two Routing processes to obtain two DigitCaps layers, respectively. These two DigitCaps layers implement the feature's coarse extraction process and the feature's fine extraction process to classify the features. There is a more gradual transition, which is beneficial to improve the classification accuracy. The dual routing is shown in Fig. 3.

Figure 3.

It should be emphasized that the importance of low-level feature capsules to high-level feature capsules can be measured by c_ij and the corresponding prediction vector. c_ij is the normalized result of b_ij. The formula is as follows: (1) cij=exp(bij)∑kexp(bij) {c_{ij}} = {{\exp \left( {{b_{ij}}} \right)} \over {\sum\nolimits_k {\exp \left( {{b_{ij}}} \right)} }}

The update of c_ij is affected by b_ij. The initial value of b_ij is 0. The c_ij normalized by the formula (1.1) is assigned the average probability value, such as: c111+c121+…+c181=1 c_{11}^1 + c_{12}^1 + \ldots + c_{18}^1 = 1 .

The update of b_ij uses the following formula: (2) bij=bij+u^ij⋅vj {b_{ij}} = {b_{ij}} + {\hat u_{ij}} \cdot {v_j}

Relying on (1) and (2) the capsule network can iteratively improve the coupling coefficient c_ij.

During training, the input of the capsule discriminator is the generated high-resolution image and the original high-resolution image of the training set. Its output is the binary classification result calculated by the second normal form of the vector output from the second DigitCaps layer, corresponding to the generated image (fake) and real image (real).

The dimension expansion of the feature vector through v¹ and v² corresponds to the rough extraction and fine extraction of the dual routing process. That allows the dual-route capsule network to deepen the network's understanding of image features in a higher dimension and enhance the discriminant network's discriminative ability to assist the generation of the network better.

Since the discriminator network is built by the capsule network as a whole, and the final actual output result is a multi-dimensional capsule vector, the impact of the capsule network's loss function needs to be considered.

L_M is the edge loss function used by the capsule network to implement parameter training, which is defined as follows: (3) LM=∑k=1KTkmax(0, m+−‖vk‖)2+λ(1−Tk)max(0, ‖vk‖−m−)2 \matrix{ {{L_M}} \hfill & = \hfill & {\sum\limits_{k = 1}^K {{T_k}\max {{\left( {0,\,{m^ + } - \left\| {{v_k}} \right\|} \right)}^2} + } } \hfill \cr {} \hfill & {} \hfill & {\lambda \left( {1 - {T_k}} \right)\max {{\left( {0,\,\left\| {{v_k}} \right\| - {m^ - }} \right)}^2}} \hfill \cr }

Among them, k is the number of classifications, and T_k is a function of the classification. T_k = 1 when and only when the classification of k appears, and T_k = 0 when it does not exist.

The loss function of the GAN discriminator is as follows: (4) max V(D, G)=Ex~pdata(x)[log D(x)]+Ez~pz(z)[log(1−D(G(z)))] \matrix{ {\max \,V\left( {D,\,G} \right)} \hfill & = \hfill & {{E_{x\sim{p_{data}}\left( x \right)}}\left[ {\log \,D\left( x \right)} \right] + } \hfill \cr {} \hfill & {} \hfill & {{E_{z\sim{p_z}\left( z \right)}}\left[ {\log \left( {1 - D\left( {G\left( z \right)} \right)} \right)} \right]} \hfill \cr }

After combining the super-resolution task with the above loss function and making improvements based on WGAN [26,27], we can define the final discriminator loss function, as in (5): (5) LD=max EIHR~ptrain(IHR)[−LM(DθD(IHR), T=1)]+EILR~pG(ILR)[−LM(DθD(GθG(ILR)), T=0)] \matrix{ {{L_D}} = {\max \,{E_{{I^{HR}}\sim{p_{train}}\left( {{I^{HR}}} \right)}}\left[ { - {L_M}\left( {{D_{{\theta _D}}}\left( {{I^{HR}}} \right),\,T = 1} \right)} \right] + } \cr \,\,\,\,\,\,\,\,\, {{E_{{I^{LR}}\sim{p_G}\left( {{I^{LR}}} \right)}}\left[ { - {L_M}\left( {{D_{{\theta _D}}}\left( {{G_{{\theta _G}}}\left( {{I^{LR}}} \right)} \right),\,T = 0} \right)} \right]} \hfill \cr }

The two parts of the formula represent the discriminative loss of the two source data, respectively. Among them, D is the discriminator network. Since the capsule network is used as the discriminator in this article, the actual output result is a vector. G is the generator network, and the output is a high-resolution image. I^HR represents the incoming discriminant network D is the training set high-resolution (HR) image, I^LR represents the incoming generation network G is the pre-processed low-resolution image (LR), The LR is processed by the G network into a high-resolution image SR and then transferred to the D network. The source is the real training set whose T value is set to 1, and the source is the generated image whose T value is set to 0.

When the generator is trained, the minimum value of the generated adversarial is calculated. Since the first half of the formula has nothing to do with the G network, the second half of the (1) is used during actual training, and the T value is set to 1, as follows: (6) LAdversarial=min EILR~pG(ILR)[LM(DθG(GθG(ILR)), T=1)] {L_{Adversarial}} = \min \,{E_{{I^{LR}}\sim{p_G}\left( {{I^{LR}}} \right)}}\left[ {{L_M}\left( {{D_{{\theta _G}}}\left( {{G_{{\theta _G}}}\left( {{I^{LR}}} \right)} \right),\,T = 1} \right)} \right]

In addition to the generator counter loss of (4), in order to improve the perception of the texture details of the generated image, this paper proposes a vector inner product loss function based on the capsule network as follows: (7) LVecoterSR=∑im(Vi(IHR)⋅Vi(GθG(ILR))−‖Vi(IHR)‖2)2 L_{Vecoter}^{SR} = \sum\limits_i^m {{{\left( {{V_i}\left( {{I^{HR}}} \right) \cdot {V_i}\left( {{G_{{\theta _G}}}\left( {{I^{LR}}} \right)} \right) - {{\left\| {{V_i}\left( {{I^{HR}}} \right)} \right\|}^2}} \right)}^2}}

V represents the vector before the compressed rectification function of the capsule inner product loss function, that is, a 16-dimensional vector from the DigitCaps layer of the capsule network. The subscript i of V represents the number of sequences to be classified, and m is the total number of classifications. There are two classifications of true and false in the experiment, and the value of i is 0 or 1.

Therefore, combining (3) and (4), the final cost function L_G of CSRGAN generator is as follows: (8) LG=LVecoterSR+10−3⋅LAdversarial {L_G} = L_{Vecoter}^{SR} + {10^{ - 3}} \cdot {L_{Adversarial}}

To make the final reconstructed pixels of the generated image closer to the HR, we need to train the main constraints to focus more on the pixels' accuracy while relatively weakening the dependence on the “creativity” of the adversarial network. The adjustment coefficient in (8) is set to 10⁻³.

The DIV2K data set is used in the experiment, which contains 800 images of the training set, 100 images of the test set, and 100 images of the verification set. Random horizontal flipping is used to enhance the data set. The bitch size of the network is set to 16. Before the experiment, it is necessary to obtain the LR image by bicubic interpolation of the HR and then use the super-resolution reconstruction algorithm to reconstruct the LR image into an SR image. The processes from LR to HR and HR to SR are all performed with a 4 times scale factor. The network's ultimate goal is to train a generator network G, which can generate a corresponding HR image from a given LR image. The experiment finally realizes the optimization of the above-mentioned mini-max problems through alternate training of D-network and G-network. The RMSProp algorithm, which is more suitable for gradient instability, is used for parameter optimization during training. In order to make the discriminator have a certain discriminative performance in the experiment, the experiment is set to start training the G network when the D network is trained 25 times. After many experiments, it is found that the training ratio of the D to G network is adjusted to 3:1. The two networks can have a better convergence effect.

The experimental content mainly focuses on the three dimensions of image texture effect, image brightness, and local features. The comparison with Bicubic, SRCNN, SRGAN, and RDN is shown in figures below:

Fig. 4, Fig. 5 and Fig. 6 respectively capture some of the 250×250 pixels, 185×185 pixels, and 100×100 pixels corresponding to the experimental results.

Figure 4.

Figure 5.

Figure 6.

In pixel-level reconstruction tasks, the texture is an important part of describing the image's details, usually in the reconstruction effect of the partial lines of the image. It can be seen from the experimental results that the reconstructed image obtained by the Bicubic method is blurred, especially in the image where the pixel difference is dense, and a lot of texture details are missing. For example, in Fig. 6, the lines on the surface of the reconstructed fruit in the image are even disconnected and discontinuous. The image reconstructed by SRCNN cannot well restore the image's texture, and there is a lack of smooth transition between local edge pixels in the high-multiple reconstruction task. For example, in Fig. 5, the image reconstructed by SRCNN does not accurately restore the straight lines of smaller wall tile targets, and some curves are reconstructed into straight lines.

In the experimental evaluation of super-resolution reconstruction, brightness is an important parameter of the similarity index of image structure. It is also the image element that human visual senses are most likely to perceive. As shown in Fig. 5, SRGAN, which uses GAN as the network structure, has a significant deviation from the label data in the restoration of the brightness space. The reconstructed building image produces a generally high brightness result. On the other hand, although the sharpness of SRGAN is greatly improved compared with SRCNN, the general grid-like texture shown in Fig. 6 is generated in the image after reconstruction, which also greatly impacts the look and feel of the image. RDN has high clarity, but it will produce excessively sharp edges in the image's local details. As shown in Fig. 5, the reconstructed image's wall lines are excessively sharp, which is not consistent with the style of ancient Roman buildings, which makes the image not realistic in appearance. CSRGAN significantly weakens the raster effect produced by the reconstruction of the SRGAN architecture and has richer local texture details, ensuring that the image is clear enough during high-magnification reconstruction.

Unlike other reconstruction methods that lack attention to the correlation between local features, CSRGAN using the capsule network is more in line with the natural distribution of the image in terms of local features. As shown in Fig. 4, we can see that in the reconstruction results of the panda's nose and surrounding hair, CSRGAN has the most natural and smooth transition and can achieve more accurate reconstruction in different parts of the hair with similar color and texture. As shown in Fig. 6, in the original image of HR, the lines and grooves on the nut's surface also have a large number of different parts of the same color but different color depths. CSRGAN also accurately achieves the reconstruction and has a look and feel closer to the real image.

In the experiment, we use PSNR as the key image evaluation index. The PSNR training results of CSRGAN based on the Set5 dataset are shown in Fig. 7.

Figure 7.

As shown in Fig. 7, CSRGAN has a relatively stable convergence process and can stably generate high-quality super-resolution images after reaching a certain number of training times.

In this experiment, the two standards of peak signal-to-noise ratio (PSNR) and structural similarity (SSIM) comprehensively evaluate the results of each experimental method as follows:

It can be seen from Table I that Bicubic's PSNR and SSIM achieved 28.39 and 0.8102, respectively, which is the lowest score for this experiment. The PSNR and SSIM parameters of SRGAN are 29.43 and 0.8477, respectively, and both indicators are weaker than SRCNN and RDN. The two parameters of CSRGAN's PSNR and SSIM's evaluation value in the experiment are compared with the leading RDN reconstruction algorithm in the field. The two parameters have increased by 0.14 and 0.0015, respectively, which shows that CSRGAN has the ability to generate more realistic super-resolution images than RDN.

Nevertheless, CSRGAN still has a lot of work worthy of in-depth research in the future, and there are also unresolved factors that negatively affect the imaging effect. In the few images in the experiment, we also found some samples with poor performance. As shown in Fig. 8 and Fig. 9 below.

Figure 8.

Figure 9.

Although the images captured in Fig. 8 and Fig. 9 are from the same set of experimental data, the image area CSRGAN shown in Example 1 exceeds the image generated by RDN in terms of the doorway's details and the overall texture.

As shown in Fig. 10 and Fig. 11, the image data is still taken from the same experimental test sample. In Fig. 10, the clarity of the reconstructed image of the trees generated by CSRGAN is far better than that of the image generated by RDN, and even the restoration of the stairs and walls in Fig. 11 gives the audience a clearer subjective feeling. However, we can observe the details of the bricks of the steps and the wall's decoration, which produces a larger deviation compared with the HR image.

Figure 10.

Figure 11.

We believe that the situation in Fig. 8 and Fig. 9 is highly likely to be caused by the fact that the capsule network has few relevant features in the training set and weak local correlation (such as the strong correlation between the overall features of the human face and the features of local facial features), and that the level of the capsule network in CSRGAN is relatively shallow and no deeper features are extracted. Improving the capsule network's performance by deepening the capsule network is a direction of our future related research. In addition to the above factors, some of the influencing factors in Fig. 10 and Fig. 11 may come from the GAN model. The selection of effective features between capsule layers is achieved through clustering algorithms, which will slow down the overall training speed during training. Therefore, how to optimize the capsule network clustering algorithm and improve its training speed in the future is also an important research direction in the research of super-resolution reconstruction of capsule network.