FNSAM: Image super-resolution using a feedback network with self-attention mechanism

2.1Feedback mechanism

For various vision tasks, the feedback mechanism has been explored by many network for carrying a notion of output to correct previous states, such as classification [21], human pose estimation [25], and other fields [26, 27]. These methods have a top-down feedback mechanism by using RNN structure. The SRFBN [22] is the first use of feedback mechanism in image SR, which concatenates the deep features of the previous iteration into the coarse feature extraction of current iteration. However, these methods directly transmit the high-level information to low-level layers, which lacks the ability to extract global features. SAM is focus on the key features for improvement of information representation. Therefore, in our method, the SAM is added to feedback network for extracting important features in feedback stage. Furthermore, the residual learning and dense skip connection are combined to build a rich image feature mapping of the network.

2.2SAM

The purpose of the attention mechanism is to focus on the details rather than the overall analysis. In recent years, spatial and channel attention mechanism has been applied to image super resolution reconstruction with promising results. Since the attention mechanism ensures that the input and output sizes remain the same, this design can be easily combined into a deep neural network for capturing dependencies in long-distance. For example, Zhang et al. proposed to use residual channel attention network to model channel correlation [28]. Dai et al. applied shifting channel attention to second-order attention enhancement [29]. Different from spatial and channel attention mechanisms, SAM exhibits promising in modeling global dependency. The SAM [30, 31] can establish the dependency relationship between the current concerned local location and all positions in the image, that is, the weighted dependency feature map of all local pixels in the image can be obtained while keeping the global features of the image. The brain MRI image exhibits non-local similarity which causes difficulty in capturing more efficient high-frequency information for super resolution reconstruction. This paper proposes to embed the SAM into the feedback network for paying attention to the weight between the pixels of the image itself, so as to enhance the expression ability of deep features in the feedback network.

3.1FNSAM framework

Generation of a feedback network calls for several conditions. Firstly, for routing the feedback information into the input at next time step, the network should be recurrent over time. Secondly, to obtain the consistency between the output and input conditions, the feedback information flow is required to have the output image feature. Thirdly, the feedback information and low-level features are combined to be the total input. RNN is a necessary condition of feedback mechanism, in which the recurrent structure is used to achieve iterative process with a LR input provided at each time step for availability of low-level information and the loss is calculated to force the network to reconstruct a SR image that guides the LR image to restore a better SR result.

The proposed FNSAM is built by RNN with time steps of T, which is shown in Fig. 1. T sub-networks are contained in our unrolled FNSAM. Each sub-network can be regarded independent at a specific time step which consists of three parts with parameters shared across time steps: the input unit, feedback unit and the output unit.

Figure 1.

The structure of our FNSAM.

The input unit is the feature extraction stage, which includes two convolution layers for extraction of initial low-level features. The first and the second layers hold 3 × 3 and 1 × 1 convolutional kernels, respectively.

Given IL⁢R as the input image at the t-th time step, the initial low-level feature xt can be obtained by

where fi⁢n denotes the function of the input unit. Then the extracted low-level feature xt is fed to feedback unit to learn more deep features.

The feedback unit is the nonlinear feature mapping stage, which consists of multiple residual dense blocks (RDBs) and a self-attention module (SAM). Multiple RDBs are used to form a better features extraction hierarchy. The SAM receives high-level information from feedback connections and extract the attention maps for further refining the low-level features. In the feedback mechanism, the SAM is used to find the principle features ht as the signal of recurrent time step t. Then, the high-level information ht at the t-th time step can be formulated as follows:

where f𝑟𝑑𝑏 represents the operation of multiple RDBs, [xt,hat-1] denotes concatenation of xt and ht.

The output unit is the image reconstruction stage, including a de-convolutional layer and a convolutional layer. The deconvolutional layer is used to upscale the high-level features and the convolutional layer with 3 × 3 sized is adopted to recover a residual image I𝑅𝑒𝑠t at the t-th time step. At the same time, the original LR image is interpolated by bilinear and combined with the residual image to obtain the SR image IS⁢Rt at the t-th time step with formulation as:

where fo and fu⁢p refer to the function of the output unit and interpolated kernel, respectively.

During each time step, L⁢1 loss is adopted to optimize the network and the loss Lt of t-th time step is calculated as:

where IGt and IS⁢Rt refer to the ground truth and the reconstruction SR image at the t-th time step, respectively. Thus, the whole loss L⁢1 can be formulated as

where Wt represents the weight of t-th loss function and the Wt is set to 1 which represents each output has equal contribution.

3.2Feedback unit

The structure of feedback unit is shown in Fig. 2. Stacking more RDBs will provide more various sizes of receptive field in each sub-network, which leads to a better hierarchy of extracted features. The HR image at the current time step is reconstructed using these features and the low-level features at the next time step are further modified. Because RDBs generate a large number of features and convolution is limited to local feature, the SAM is cooperated with feedback mechanism which can capture global dependencies and improve the accuracy of features.

Figure 2.

The proposed feedback unit with RDB and self-attention module.

The high-level features ht-1 at t-1-th time step are served as the input of the SAM to generate important features ht. The RDBs simultaneously receives the low-level features xt of the input unit and the output ht of the SAM at t-th time step. Supposing there are r RDBs, the output ur of the r-th RDB can be obtained by

where f𝑟𝑑𝑏r refers to the operations of the r-th RDB, [xt,hat-1] represents the concatenation of xt and hat-1.

The architecture of RDB is shown in Fig. 2b which is the same as the one in [16]. It contains dense connected layers, feature fusion and local residual learning. For each convolution layer, firstly, multiple trainable kernels are used to convolve the feature map of previous layer for extraction of the more informative features and generation of this layer feature maps. Then the activation function is used to non-linearize the convolutional result to yield more complex characteristics. The output ulr of the l-th convolutional layer of the r-th RDB is described as:

where f⁢(⋅) is the activation function, wlr denotes the weights of the l-th convolutional layer of the r-th RDB, [ur-1,u1r,⋯,ul-1r] represents the operation of concatenation. Since many features are generated by RDB, converting features into low-level features requires a lot of parameters and computational burden. Therefore, a 1 × 1 convolutional kernel is used to compress the features which can be formulated as

where f𝑐𝑜𝑛𝑣r denotes the function of the 1 × 1 convolutional layer in the r-th RDB and L is the number of layers in each RDB. In order to alleviate the problem of network gradient vanishing, local residual learning is used. Therefore, the output of the r-th RDB is calculated as:

3.3Self-attention module

The self-attention module is divided into three branches, which decompose the input data into three components: θ, ϕ and g. θ and ϕ are used to study the global dependency relationship between the characteristics. g is used to extract the deep features of the input feature maps. Then the global dependency and deep features are fused to enhance the capacity of characteristics graph representation.

For the feature maps x, the attention map s⁢(x) can be calculated as:

where θ⁢(x)=Wθ⁢x, ϕ⁢(x)=Wϕ⁢x, θ⁢(x)T⁢ϕ⁢(x) represents the dot product operation of the matrix, T refers to the transpose of the matrix, and the weights between two positions of input feature map x are computed by performing the softmax on each row of the result. The attention map s⁢(x) is built using these weights for description of the similarity between θ⁢(x) and ϕ⁢(x).

Then, the attention map s⁢(x) and feature map g⁢(x) are multiplied, which is added with the input feature maps x via the skip connection to get the output self-attention feature maps O⁢(x). The calculation process is

where g⁢(x)=Wg⁢x. In the above formulation, Wθ, Wϕ, Wg and Ws are the learned weight matrices, which are implemented as 1 × 1 convolutions.

It can be seen from the above description that the SAM model can focus on the pixel data of the whole area of the image. When integrated with deep features, the SAM can easily capture the main features in the whole spatial domain and enrich the context dependence according to the self-attention features.

4.1Datasets and training settings

Medical images used in this paper are from the brain MRI dataset in The Cancer Imaging Archive (TCIA) [32] public database. We randomly select 80 representative images as the training dataset and 44 images as the test dataset which is divided into Test1 (8 images), Test2 (14 images) and Test3 (22 images). Data enhancement is performed on the selected training set using rotation and flips [33]. Specifically, the original images are rotated by 90∘, 180∘, 270∘ and reduced to 0.6, 0.7, 0.8 times of the original size. Then, we obtain 16-fold images, for a total of 1280 training images.

The first layer has 256 convolutional kernels and 64 RDBs. Because our purpose is to reconstruct MRI images, the number of input and output channels can be set to 1. Same as RDN [16], per RDB has 8 convolutional layers. The kernel size, stride, and padding of all convolutional layers in RDBs and the convolutional layer in output unit are set to 3, 1, and 1, respectively which can keep size fixed. A PReLU [35] activation function is used to follow all convolutional layers except the last convolutional layer of each RDB and the sub-network. The de-convolutional layer is set for each upscale factor as shown in Table 1.

The training image is divided into 48 × 48 small blocks and the step size 48 is used to increase the image receptive field. For self-attention module, we use 1 × 1 kernels with stride 1 and pad 1. Each weighting layer has 64 convolutional kernels of size 3 × 3. The minimum batch size is 16. For × 2 and × 3 scale factor, the LR input patch size is 40. For × 4 scale factor, the LR input patch size is 32. In addition, the Adam algorithm [34] is applied for gradient optimization with the initial learning rate 0.0001. We make the learning rate halved for each 2 × 105 iterations. Experiment is implemented on a workstation with a GTX 1080Ti GPU and the memory of 11 GB under the software environment of Ubuntu 16.04 and MATLAB. To qualitatively evaluate the effectiveness of different methods, two image quality evaluation indexes, peak signal-to-noise ratio (PSNR) [35] and structural similarity index measure (SSIM) [36] are adopted in this paper. The higher PSNR and SSIM indicate the more reasonable SR quality.

4.2Ablation analysis

To verify the effectiveness of feedback network and SAM, ablation analysis is performed. As shown in Table 2, the results of PSNR and SSIM on Test3 with upscale 4 reveal the advantage of feedback network and SAM. In the second group, it can be seen that using feedback network would perform better than the first group. This is mainly because RNN based feedback mechanism can reduce the errors occurring at each preceding time step. This further shows that the output feature map containing advanced information in the t-th iteration will prompt the previous layers to generate better representations in subsequent iterations. In the third group, the SAM is added in feedback network which can improve the performance. It is because the global weight distribution learned chooses more useful features to be the signal of feedback.

4.3R and T analysis

The RDB number (denoted as R) and the time step number (denoted as T) are studied. Firstly, we construct various networks FNSAM-S (R = 4), FNSAM-M (R = 5), FNSAM-L (R = 6) and FNSAM-L+ (R = 7) by fixing T to 4. The performance of different number of RDBs on Test3 is illustrated in Fig. 3. Generally speaking, performance will get better as the network depth increases. However, it can be found that the red line (R = 5) is higher than the blue line (R = 6) and the yellow line (R = 7). The average results of on 4 × enlargements are shown in Table 3. The network performance is basically saturated when the R is larger than 5. Considering that the difference of PSNR between FNSAM-M and FNSAM-L networks is only 0.02 dB, FNSAM-M (R = 5) is chosen for the time step analysis.

Figure 3.

The convergence analysis of R.

Figure 4.

The convergence analysis of T.

Under identical setting of R = 5, we set T = 2, 3, 4, and 5, respectively. Figure 4 shows the convergence of T on Test3 with scaling factor × 4. Table 4 shows the specific numerical comparison. According to the two groups results, we can observe that the quality of reconstruction keeps rising with an increase of T except for a slight decrease when T = 5. This indicates the feedback mechanism contributes to the flow of information across time. Therefore, we choose T = 4 and R = 5 to balance the reconstruction quality and the computational complexity.

4.4Comparison with the state-of-the-arts

The proposed SR algorithm is compared with the current classical algorithms including ScSR [9], SRCNN [11], VDSR [12], EDSR [14], D-DBPN [15], RDN [16] and SRFBN [22] methods. Tables 5 and 6 show the average PSNR and SSIM for scale factor × 2, × 3 and × 4 on Test1, Test2, Test3, respectively. Compared with the classical algorithms, our method is obviously higher than those of all the compared methods both in PSNR and SSIM. For different up-sampling factors, the PSNR and SSIM values of our method are slightly higher than the SRFBN in some test images. For scale factor × 4 on Test3, the PSNR produced by the proposed method is 2.59 dB, 1.96 dB, 1.02 dB, 0.62 dB, 0.84 dB, 0.35 dB and 0.16 dB higher than that of ScSR, SRCNN, VDSR, EDSR, D-DBPN, RDN and SRFBN, respectively, while the increase of SSIM is 5.81%, 3.85%, 1.43%, 0.99%, 1.33%, 0.49% and 0.13%, respectively.

Figure 5.

Visual results of different SR approaches for different scale factors.

Figure 5 gives the visualization of SR results using different methods. As shown in these figures, details are lost obvious in SR reconstruction produce by the bicubic interpolation. The proposed FNSAM alleviates the distortions and generates more accurate details in SR images compared with other methods. For example, the pattern in Fig. 5a is supposed to a vertical line. Affected by the surrounding stripe on the original image, other SR methods tend to reconstruct the curved pattern, and the proposed FNSAM can reconstruct the pattern correctly. In Fig. 5b and c, the proposed FNSAM can reconstruct the more obvious detail than others. In other words, our methods can recover clearer edges and more textures which are close to HR medical image than the comparison methods.

4.5Network performance and parameters

Figure 6.

Performance and number of parameters.

Figure 6 shows the comparison between the performance and parameters of some advanced SR methods and the proposed method on Test3 with upscale 4. As can be seen that parameters of SRCNN, VDSR and SRFBN methods are fewer but with lower PSNR than our proposed model. Our FNSAM achieves better performance and fewer parameters than D-DBPN, RDN and EDSR methods. This is attributed to the ability of RNN structure to decrease the parameters of the network.