An Image Restoration Method and Device Based on Distortion Estimation

Abstract

The present invention discloses an image restoration method and apparatus based on distortion estimation. The method includes: extracting image features from an original image to obtain an image feature map; inputting the image feature map into a distortion factor estimation module to obtain a distortion vector composed of one or more factors causing the distortion of the original image; inputting the distortion vector and the image feature map into an image reconstruction module, and the image reconstruction module removes the corresponding distortion and performs image reconstruction to obtain a residual image corresponding to the original image; based on the residual image, a restored image corresponding to the original image is obtained. The method can process one or more types of mixed distortion, and has strong applicability and generality.

Description

Technical Field

[0001] The present invention relates to the field of image processing, and particularly to an image restoration method and apparatus based on distortion estimation. Background Art

[0002] Image restoration is a very important processing technology in the field of image processing, and its purpose is to reconstruct a high-quality clear image from a distorted image. At present, image restoration technology has been widely applied in fields such as aerospace, military and public security, biomedicine, and industrial production, and has broad application prospects.

[0003] Most current image restoration methods are for single-distortion image restoration problems, such as image denoising, image deblurring, and JPEG compression distortion removal. However, in actual scenarios, images mostly experience multiple mixed distortions, and the image restoration methods designed for single-distortion images are no longer applicable.

[0004] In recent years, many researchers have studied image restoration methods for mixed-distortion images. The current restoration methods can be roughly divided into two categories: one is an image restoration method that can handle specific mixed distortions, and the other is an image restoration method for mixed-distortion images that can handle unknown distortion types.

[0005] Image restoration methods capable of handling specific mixed distortions require predefined distortion types and then adopt different strategies to restore specific mixed-distortion images. In 2018, Yu K et al. proposed an image restoration algorithm based on deep reinforcement learning in "Crafting a Toolchain for Image Restoration by Deep Reinforcement Learning", which gradually restored images with Gaussian blur, Gaussian noise, and JPEG compression distortion simultaneously. However, the restoration accuracy of this method for mixed-distortion images is limited, and at the same time, the efficiency is low because each small-scale convolutional network needs to be pre-trained. Therefore, in 2019, Suganuma M et al. proposed a neural network that performs multiple operations in parallel and uses an attention mechanism to adaptively select operations for mixed distortions in "Attention-based Adaptive Selection of Operations for Image Restoration in the Presence of Unknown Combined Distortion". However, the restoration accuracy of this method for mixed-distortion images still needs to be improved, and at the same time, a large number of different operations have also increased the computational parameters of this method. In addition, the designs of the above two methods ignore the interference between different distortion types. In 2020, Li X et al. proposed a hybrid-distortion image restoration network based on disentangled feature representation in "Learning Disentangled Feature Representation for Hybrid-distorted Image Restoration". This network reduces the interference between different distortions through feature disentanglement and can solve the restoration problem of mixed-distortion images containing Gaussian noise, Gaussian blur, and JPEG compression distortion simultaneously, and the restoration quality is further improved. In 2019, Qian G et al. proposed the Trinity Enhancement Network TENet in "Trinity of Pixel Enhancement: a Joint Solution for Demosaicking, Denoising and Super-Resolution". For mixed-distortion images with low resolution, noise, and mosaic simultaneously, this network adopts an image pipeline processing and joint learning strategy, first performs image denoising, then performs image super-resolution, and finally performs image demosaicking.In addition, in 2019, Li Jiafeng et al. disclosed a high-quality reconstruction method for compound degraded images based on generative adversarial networks in the patent "A High-quality Reconstruction Method for Compound Degraded Images Based on Generative Adversarial Networks". This method uses the game between the generative network and the discriminative network to restore low-quality images that simultaneously contain multiple degradation problems such as haze, system noise, low illumination, and compression distortion. The above methods can only process specific mixed distortion images and cannot restore other types of mixed distortion images and single distortion images.

[0006] In actual scenarios, images generally undergo mixed distortions of unknown types. In recent years, some researchers have begun to study restoration methods for mixed distortion images that can handle unknown distortion types. In 2020, Shan Chaowei proposed a hybrid multi-distortion enhancement algorithm based on dual prior learning in "Research on Task-driven Deep Learning Image Enhancement". The experimental results show that the proposed algorithm can more effectively process mixed distortion images with unknown distortion types and degrees compared to the original deep image prior. In 2020, Hou Shujuan et al. disclosed a general image restoration method based on neural networks in the patent "A General Image Restoration Method Based on Neural Networks". For mixed distortion images with unknown distortion types, this method uses SVM to judge the distortion type and then performs corresponding distortion repair according to the specified pipeline sequence to obtain the restored image.

[0007] However, the existing image restoration methods for images experiencing multiple distortions have the following disadvantages: (1) Although some methods can process images with three types of mixed distortions, these methods cannot process the corresponding single distortion images, and the compatibility of the methods is poor. (2) Some methods can achieve the restoration of unknown mixed distortions. For example, the hybrid multi-distortion enhancement algorithm based on dual prior learning in the literature "Research on Task-driven Deep Learning Image Enhancement" and the method in the patent "A General Image Restoration Method Based on Neural Networks". The former uses a semi-supervised learning method, and the restoration effect of the image is poor; the latter's distortion detection module uses SVM to judge the distortion type, does not consider the impact of the distortion degree of the image on image restoration, and the distortion repair module uses a pipeline strategy and does not consider the mutual interference between distortion types. When removing one distortion, some new distortions will be introduced, resulting in poor quality of the restored image. Summary of the Invention

[0008] In view of this, the present invention provides an image restoration method and device based on distortion estimation, which can solve the general image restoration technical problems. The present invention can not only handle various mixed distortions but also handle single distortions, with strong applicability and generality and good restoration effects.

[0009] To solve the above technical problems, the present invention is implemented as follows.

[0010] An image restoration method based on distortion estimation, comprising:

[0011] Step S1: Obtain an original image, and extract image features from the original image to obtain an image feature map;

[0012] Step S2: Input the image feature map into a distortion factor estimation module to estimate one or more factors causing the distortion of the original image, and obtain a distortion vector composed of the one or more factors causing the distortion of the original image; the distortion factor estimation module includes a convolutional layer, a global average pooling layer, and a fully connected layer, and each factor corresponds to a type of image distortion;

[0013] Step S3: Input the distortion vector and the image feature map into an image reconstruction module, and the image reconstruction module removes the corresponding distortion and performs image reconstruction to obtain a residual image corresponding to the original image; based on the residual image, obtain a restored image corresponding to the original image; wherein, the image reconstruction module includes an attention convolutional block for removing distortion and performing image reconstruction.

[0014] Preferably, the factor estimation module includes a first convolutional layer, a global average pooling layer, a first fully connected layer, a first ReLU activation function layer, a second fully connected layer, a second ReLU activation function layer, a third fully connected layer, and a Sigmoid activation function layer connected in sequence. Among them, the factor estimation module inputs a distorted image, and after being processed by the factor estimation module, outputs a distortion vector with the number of elements being k. The distortion factors are the respective elements in the vector, realizing the estimation of one or more factors causing the distortion of the original image.

[0015] Preferably, the image reconstruction module includes n attention convolutional blocks, a second convolutional layer, a third ReLU activation function layer, and a third convolutional layer connected in sequence.

[0016] Preferably, the attention convolutional block is used to remove distortion and perform image reconstruction; the attention convolutional block includes a reconstruction branch and an adjustment branch;

[0017] The reconstruction branch is used to receive the image feature map or the feature map output by the previous attention convolutional block of the current attention convolutional block; if the current attention convolutional block is the first-level attention convolutional block, the input is the image feature map; if the current attention convolutional block is a non-first-level attention convolutional block, the input is the feature map output by the previous attention convolutional block of the current attention convolutional block. The reconstruction branch includes an attention sub-branch and a non-attention sub-branch. The attention sub-branch is used to receive the image feature map or the feature map output by the previous attention convolutional block of the current attention convolutional block; if the current attention convolutional block is the first-level attention convolutional block, the input is the image feature map; if the current attention convolutional block is a non-first-level attention convolutional block, the input is the feature map output by the previous attention convolutional block of the current attention convolutional block. The attention sub-branch includes a fourth convolutional layer, a fourth ReLU activation function layer, a fifth convolutional layer, a spatial attention module, and a channel attention module connected in sequence. The attention sub-branch is used to remove distortion and reconstruct the image, and output the attention sub-branch residual feature map. The spatial attention module is used to enhance the image texture and detail information in the space, and suppress the interference information of distortion such as noise in the space. The channel attention module is used to enhance the channels containing image texture and detail information, and suppress the channels containing interference information of distortion such as noise, and output the attention sub-branch residual feature map; wherein, the attention sub-branch is used to multiply with the output of the first distortion factor attention module of the adjustment branch to obtain a first fusion result;

[0018] The non-attention sub-branch is used to receive the image feature map or the feature map output by the previous attention convolutional block of the current attention convolutional block; if the current attention convolutional block is the first-level attention convolutional block, the input is the image feature map; if the current attention convolutional block is a non-first-level attention convolutional block, the input is the feature map output by the previous attention convolutional block of the current attention convolutional block. The non-attention sub-branch includes a sixth convolutional layer, a fifth ReLU activation function layer, and a seventh convolutional layer connected in sequence. The non-attention sub-branch is used to supplement the feature information ignored by the attention sub-branch, and output the non-attention sub-branch residual feature map. The non-attention sub-branch residual feature map is used to multiply with the output of the second distortion factor attention module of the adjustment branch to obtain a second fusion result;

[0019] The adjustment branch includes a first distortion factor attention module and a second distortion factor attention module. The input of the adjustment branch is the distortion vector. The first distortion factor attention module is used to generate a weight vector for controlling the attention sub-branch of the reconstruction branch, multiply the output with the output of the attention sub-branch to obtain a first fusion result. The second distortion factor attention module is used to generate a weight vector for controlling the non-attention sub-branch of the reconstruction branch, and fuse the output with the output of the non-attention sub-branch to obtain a second fusion result.

[0020] Add the first fusion result and the second fusion result to obtain a third fusion result. Add the third fusion result to the image feature map to obtain a processed image feature map, which is used as the output of the attention convolution block.

[0021] Preferably, based on the residual image, obtaining a restored image corresponding to the original image includes: adding the obtained residual image and the original image pixel by pixel to obtain the restored image corresponding to the original image, and outputting the restored image.

[0022] An image restoration device based on distortion estimation provided by the present invention, the device includes:

[0023] An image feature map acquisition module: configured to acquire an original image, and extract image features from the original image to obtain an image feature map;

[0024] A distortion estimation module: configured to input the image feature map into a distortion factor estimation module, estimate one or more factors causing the distortion of the original image, and obtain a distortion vector composed of the one or more factors causing the distortion of the original image; the distortion factor estimation module includes a convolutional layer, a global average pooling layer, and a fully connected layer, and each factor corresponds to a type of image distortion.

[0025] An image reconstruction module: configured to input the distortion vector and the image feature map into the image reconstruction module, remove the corresponding distortion by the image reconstruction module, and perform image reconstruction to obtain a residual image corresponding to the original image; based on the residual image, obtain a restored image corresponding to the original image; wherein, the image reconstruction module includes an attention convolution block for removing distortion and performing image reconstruction.

[0026] A computer-readable storage medium, in which multiple instructions are stored; the multiple instructions are used to be loaded and executed by a processor to perform the method as described above.

[0027] An electronic device, characterized in that the electronic device includes:

[0028] A processor for executing multiple instructions;

[0029] A memory for storing a plurality of instructions;

[0030] The plurality of instructions are used to be stored by the memory and loaded and executed by the processor to implement the method as described above.

[0031] The present invention proposes a universal image restoration method based on distortion estimation. The method processes mixed distorted images with unknown distortion type and degree in actual scenes. First, features are extracted from the input image to obtain an image feature map. Then, the image feature map is used to estimate the distortion type and degree in the image to obtain a distortion vector. Finally, based on the image feature map and the distortion vector, an image reconstruction module constructed by an attention convolution block is used to obtain the restored image.

[0032] Beneficial effects:

[0033] (1) The distortion factor estimation module of the present invention can not only determine the type of distortion contained in the image to be processed, but also estimate the corresponding degree of distortion, and the error of the estimation result is small.

[0034] (2) The image reconstruction module of the present invention utilizes the distortion factor attention to convert the distortion vector into weight coefficients that control the attention and non-attention sub-branches, thereby adjusting the image reconstruction process to adapt to the restoration of images with different distortion types and degrees. The design of this image restoration module considers the impact of multiple distortions. Compared with existing pipelined methods that remove distortions one by one, it avoids the introduction of new distortions and achieves better restoration results.

[0035] (3) The method of the present invention can process one or more mixed distortions and has strong applicability and versatility.

[0036] (4) The method of the present invention has high accuracy in image restoration. BRIEF DESCRIPTION OF THE DRAWINGS

[0037] Figure 1 A schematic flow chart of the image restoration method based on distortion estimation provided by the present invention;

[0038] Figure 2 A schematic diagram of the architecture of the image restoration method based on distortion estimation provided by the present invention;

[0039] Figure 3 Schematic diagram of the attention convolution block structure provided by the present invention;

[0040] Figure 4 This is a schematic diagram of the structure of the distortion factor attention module provided by the present invention;

[0041] Figure 5Schematic structural diagram of the image restoration device based on distortion estimation provided by the present invention. Detailed implementation manners

[0042] As Figure 1 - Figure 2 shown, the present invention proposes an image restoration method based on distortion estimation, including the following steps:

[0043] Step S1: Obtain an original image, and extract image features from the original image to obtain an image feature map;

[0044] Step S2: Input the image feature map into a distortion factor estimation module to estimate one or more factors causing distortion of the original image, and obtain a distortion vector composed of the one or more factors causing distortion of the original image; the distortion factor estimation module includes a convolutional layer, a global average pooling layer, and a fully connected layer, and each factor corresponds to a type of image distortion;

[0045] Step S3: Input the distortion vector and the image feature map into an image reconstruction module, and the image reconstruction module removes the corresponding distortion and performs image reconstruction to obtain a residual image corresponding to the original image; based on the residual image, obtain a restored image corresponding to the original image; wherein, the image reconstruction module includes an attention convolutional block for removing distortion and performing image reconstruction.

[0046] The image restoration method based on distortion estimation of the present invention first extracts features from the input image to obtain an image feature map; then uses the image feature map to estimate the type and degree of distortion in the image to obtain a distortion vector; finally, uses the image feature map and the distortion vector to obtain a restored image through the processing of an image reconstruction module constructed by attention convolutional blocks.

[0047] In the step S1, where:

[0048] The original image refers to an image collected by an image acquisition device under different scenarios, and there will be mixed distortions of unknown types and unknown degrees in these original images, such as fog, Gaussian blur, Gaussian noise, JPEG compression distortion, etc.

[0049] In this embodiment, an image feature extraction module extracts image features from the original image. The feature extraction module is composed of 1 convolutional layer and m convolutional blocks, where 3 < m < 10; after the original image undergoes feature extraction by the feature extraction module, an image feature map with dimensions of W×H×C is obtained for subsequent module use, and W, H, and C respectively represent the width, height, and number of channels of the image feature map.

[0050] In the step S2, where:

[0051] The factor estimation module includes a first convolutional layer, a global average pooling layer, a first fully connected layer, a first ReLU activation function layer, a second fully connected layer, a second ReLU activation function layer, a third fully connected layer, and a Sigmoid activation function layer, which are connected in sequence. Among them, the factor estimation module inputs a distorted image. After being processed by the factor estimation module, the distorted image outputs a distortion vector with k elements. The distortion factors are the respective elements in the vector, realizing the estimation of one or more factors causing the distortion of the original image.

[0052] Each of the factors corresponds to a type and can represent the degree of distortion caused by the distortion reason corresponding to that type. The larger the factor value, the greater the degree of distortion. When the factor value is 0, it means that the original image does not include the distortion caused by the distortion reason corresponding to the type of this factor.

[0053] The step S3: Input the distortion vector and the image feature map into the image reconstruction module. The image reconstruction module removes the corresponding distortion and performs image reconstruction to obtain the residual image corresponding to the original image; based on the residual image, obtain the restored image corresponding to the original image; among them, the image reconstruction module includes an attention convolutional block for removing distortion and performing image reconstruction, where:

[0054] The image reconstruction module includes n attention convolutional blocks, a second convolutional layer, a third ReLU activation function layer, and a third convolutional layer, which are connected in sequence, where 10 < n < 30. Input the distortion vector and the image feature map into the image reconstruction module. The image reconstruction module removes the corresponding distortion and performs image reconstruction to obtain the residual image corresponding to the original image; based on the residual image, obtain the restored image corresponding to the original image.

[0055] Such as Figure 3 - Figure 4As shown, the attention convolution block is used to remove distortion and reconstruct the image. The attention convolution block includes a reconstruction branch and an adjustment branch. The reconstruction branch is used to receive the image feature map or the feature map output by the previous attention convolution block of the current attention convolution block. If the current attention convolution block is the first-level attention convolution block, the input is the image feature map. If the current attention convolution block is a non-first-level attention convolution block, the input is the feature map output by the previous attention convolution block of the current attention convolution block. The reconstruction branch includes an attention sub-branch and a non-attention sub-branch. The attention sub-branch includes a fourth convolutional layer, a fourth ReLU activation function layer, a fifth convolutional layer, a spatial attention module, and a channel attention module connected in sequence. The attention sub-branch is used to remove distortion and reconstruct the image, and outputs an attention sub-branch residual feature map. The spatial attention module is used to enhance the image texture and detail information in the space, and suppress the interference information such as noise distortion in the space. The channel attention module is used to enhance the channels containing image texture and detail information, and suppress the channels containing interference information such as noise distortion, and outputs an attention sub-branch residual feature map. Among them, the attention sub-branch is used to multiply with the output of the first distortion factor attention module of the adjustment branch to obtain a first fusion result. The non-attention sub-branch is used to receive the image feature map or the feature map output by the previous attention convolution block of the current attention convolution block. If the current attention convolution block is the first-level attention convolution block, the input is the image feature map. If the current attention convolution block is a non-first-level attention convolution block, the input is the feature map output by the previous attention convolution block of the current attention convolution block. The non-attention sub-branch includes a sixth convolutional layer, a fifth ReLU activation function layer, and a seventh convolutional layer connected in sequence. The non-attention sub-branch is used to supplement the feature information ignored by the attention sub-branch, and outputs a non-attention sub-branch residual feature map. The non-attention sub-branch residual feature map is used to multiply with the output of the second distortion factor attention module of the adjustment branch to obtain a second fusion result. The adjustment branch includes a first distortion factor attention module and a second distortion factor attention module. The input of the adjustment branch is the distortion vector. The first distortion factor attention module is used to generate a weight vector for controlling the attention sub-branch of the reconstruction branch, and multiply the output with the output of the attention sub-branch to obtain a first fusion result. The second distortion factor attention module is used to generate a weight vector for controlling the non-attention sub-branch of the reconstruction branch, and fuse the output with the output of the non-attention sub-branch to obtain a second fusion result. Add the first fusion result and the second fusion result to obtain a third fusion result, and add the third fusion result to the image feature map to obtain a processed image feature map as the output of the attention convolution block.

[0056] The corresponding distortion is removed by the image reconstruction module, and image reconstruction is performed to obtain the residual image corresponding to the original image.

[0057] Further, the outputs of the first distortion factor attention module and the second distortion factor attention module respectively represent the weights of the image feature maps generated by controlling the non-attention sub-branch and the attention sub-branch. The distortion vector assigns weight coefficients to the attention sub-branch and the non-attention sub-branch of each attention convolution block reconstruction branch through the distortion factor attention module, so as to adapt to the reconstruction of images with different distortion types and degrees, and obtain the residual image corresponding to the original image.

[0058] Both the first distortion factor attention module and the second distortion factor attention module are distortion factor attention modules. The distortion factor attention module includes a third fully connected layer and a second Sigmoid activation function. The input of the distortion factor attention module is the distortion vector, and the output is the weight. The output of the first distortion factor attention module is the first weight α1, and the output of the second distortion factor attention module is the second weight α2.

[0059] In this embodiment, it is defined that X is the input of the attention convolution block and Y is the output of the attention convolution block. Then the attention convolution block is expressed as Y = f1(X)×α1 + f2(X)×α2 + X, where f1(X) and f2(X) respectively represent the residual feature maps generated by the attention sub-branch and the non-attention sub-branch, and α1 and α2 respectively represent the weight vectors controlling the attention sub-branch of the reconstruction branch and the weight vector controlling the non-attention sub-branch of the reconstruction branch. Through the weight vectors α1 and α2, the output Y can be continuously changed from X to f1(X) + f2(X) + X; when α1 = α2 = 0, the residual part will be skipped and not processed. α1 and α2 are respectively generated by two distortion factor attention modules. The distortion factor attention module consists of 1 fully connected layer and 1 Sigmoid function. The specific structure diagram is as Figure 4 shown. Both the attention sub-branch and the non-attention sub-branch in each attention convolution block are regarded as a functional unit, with a total of 2n functional units. By adjusting α1 and α2 of each attention convolution block, weights are assigned to each functional unit, so that the image reconstruction module can process different types and degrees of distortion.

[0060] Based on the residual image, a restored image corresponding to the original image is obtained, including: adding the obtained residual image and the original image pixel by pixel to obtain the restored image corresponding to the original image, and outputting the restored image.

[0061] Further, during the training process, the feature extraction module, the distortion factor estimation module, and the image reconstruction module are jointly trained. The feature extraction module is used to obtain the original image and extract image features from the original image to obtain an image feature map. Define the restored image as The undistorted image is The estimated distortion vector is The true distortion vector is Then the loss function is

[0062]

[0063] where N is the batch size during training, is the L1 loss between the restored image corresponding to the original image and the undistorted image corresponding to the original image, is the L1 loss between the estimated distortion vector and the true distortion vector. i is the serial number of the image and the corresponding distortion vector, and the image is the restored image or the undistorted image, and the numbers of the restored image and the undistorted image are the same. The training iterator selects AdamW. This loss function reflects the weighted sum of the L1 loss between the restored image corresponding to the original image and the undistorted image corresponding to the original image and the L1 loss between the estimated distortion vector and the true distortion vector.

[0064] To better illustrate the purpose and advantages of the present invention, specific embodiments are described below, and further explanations are made in conjunction with the drawings, experiments, and data analysis.

[0065] Assume that the obtained image contains one, two, three, or no distortion among Gaussian blur, Gaussian noise, and JPEG compression distortion.

[0066] (1) Obtain the original image.

[0067] The original image refers to the images captured by image acquisition devices under different scenarios, and there will be mixed distortions of unknown types and degrees in these images. Considering that there is currently no dataset that meets the requirements, the present invention designs and generates a dataset containing three types of mixed distortions (Gaussian blur, Gaussian noise, and JPEG compression distortion), where Gaussian blur, Gaussian noise, and JPEG compression distortion are sequentially added to the clear image with different distortion degrees. The standard deviation σ1 of Gaussian blur ∈ [0, 5], the standard deviation σ2 of Gaussian noise ∈ [0, 50], and the quality factor q of JPEG compression distortion ∈ [10, 100]. This dataset contains 8 types of images, namely: images containing all three types of mixed distortions of Gaussian blur, Gaussian noise, and JPEG compression distortion, a total of 1 type; distorted images containing any two of the above three distortions, a total of 3 types; distorted images containing any one of the above three distortions, a total of 3 types; images without distortion, a total of 1 type.

[0068] (2) Through the feature extraction module, feature extraction is performed on the original image to obtain an image feature map.

[0069] The feature extraction module consists of 1 3×3 convolutional layer and m (3 < m < 10) convolutional blocks, and the specific structure is as Figure 1 shown. Each convolutional block contains two 3×3 convolutional layers, and the ReLU function is used for activation between the two convolutional layers. The feature extraction module performs feature extraction on the input original image to obtain an image feature map of W×H×C, which is convenient for subsequent distortion factor estimation module and image reconstruction module to use, where W, H, and C respectively represent the width, height, and number of channels of the image feature map.

[0070] (3) Input the image feature map into the distortion factor estimation module to estimate the distortion factor of each distortion in the image and obtain a distortion vector.

[0071] The present invention uses the distortion factor qf ∈ [0, 1] to represent the distortion degree of each distortion in the image. The larger qf is, the greater the distortion degree. When qf = 0, it means that the image does not contain this distortion. As Figure 2 shown, the distortion factor estimation module first uses a 3×3 convolutional layer with a stride of 2 as a downsampling layer to downsample the input image feature map; then uses a global average pooling layer to compress the downsampled feature map in the spatial dimension to obtain a one-dimensional vector; secondly, there are 3 fully connected layers, and the ReLU function is used for activation between every 2 fully connected layers to estimate the distortion factors qf1, qf2, and qf3 of Gaussian blur, Gaussian noise, and JPEG compression distortion contained in the image; finally, the Sigmoid function is used to map the distortion factor to [0, 1], and finally a distortion vector QF = [qf1, qf2, qf3] containing three distortion factors is obtained, where

[0072] (4) Input the image feature map and the distortion vector into the image reconstruction module for corresponding distortion removal and image reconstruction to obtain a residual image.

[0073] The image reconstruction module consists of n (10 < n < 30) attention convolution blocks, 2 3×3 convolution layers, and 1 ReLU activation function, and its specific structure is as Figure 2 shown. The image reconstruction module takes the image feature map and the distortion vector as inputs. The distortion vector adjusts the attention convolution blocks to adapt to the reconstruction of images with different distortion types and degrees, and a residual image is obtained.

[0074] Next, the attention convolution block will be introduced in detail.

[0075] As Figure 3 shown, the attention convolution block processes the input feature map into two branches: an upper branch and a lower branch. The upper and lower branches are the attention branch and the non-attention branch respectively. The attention branch consists of 2 3×3 convolution layers, 1 ReLU activation function layer, a spatial attention mechanism, and a channel attention mechanism; the non-attention branch consists of 2 3×3 convolution layers and 1 ReLU activation function layer. Define X as the input feature map and Y as the output feature map. Then the attention convolution block can be expressed as Y = f1(X)×α1 + f2(X)×α2 + X, where f1(X) and f2(X) represent the residual feature maps generated by the attention branch and the non-attention branch respectively, and α1 and α2 represent the adjustment variables that control the feature maps generated by the two branches. By adjusting the variables α1 and α2, the output Y can continuously vary from X to f1(X) + f2(X) + X; when α1 = α2 = 0, the residual part will be skipped and not processed. α1 and α2 are generated by two distortion factor attention modules respectively. The distortion factor attention module consists of 1 fully connected layer and 1 Sigmoid function, and its specific structure diagram is as Figure 4 shown. Both the attention branch and the non-attention branch in each attention convolution block are regarded as a functional unit, with a total of 2n functional units. By adjusting α1 and α2 of each attention convolution block, weighted coefficients are assigned to each functional unit, enabling the network to process different types and degrees of distortion.

[0076] (5) Output the restored image.

[0077] Sum the obtained residual image and the original image pixel by pixel to output the restored image.

[0078] The above general image restoration method based on distortion estimation requires joint training of the feature extraction module, the distortion factor estimation module, and the image reconstruction module. Define the restored image as the clear image as the estimated distortion vector as The true distortion vector is Then the loss function is

[0079]

[0080] where N is the batch size during training. The training iterator selects AdamW.

[0081] The present invention also provides an image restoration device based on distortion estimation, as Figure 5 shown. The device includes:

[0082] An image feature map acquisition module: configured to acquire an original image and extract image features from the original image to obtain an image feature map;

[0083] A distortion estimation module: configured to input the image feature map into a distortion factor estimation module to estimate one or more factors causing the distortion of the original image, and obtain a distortion vector composed of the one or more factors causing the distortion of the original image; the distortion factor estimation module includes a convolutional layer, a global average pooling layer, and a fully connected layer, and each factor corresponds to a type of image distortion;

[0084] An image reconstruction module: configured to input the distortion vector and the image feature map into the image reconstruction module, and the image reconstruction module removes the corresponding distortion and performs image reconstruction to obtain a residual image corresponding to the original image; based on the residual image, a restored image corresponding to the original image is obtained; wherein, the image reconstruction module includes an attention convolutional block for removing distortion and performing image reconstruction.

[0085] The above specific embodiments only describe the design principle of the present invention. The shapes and names of the components in this description can be different and are not limited. Therefore, those skilled in the art of the present invention can modify or equivalently replace the technical solutions recorded in the foregoing embodiments; and these modifications and replacements do not depart from the spirit and technical solutions of the present invention, and shall all fall within the protection scope of the present invention.

Claims

1. An image restoration method based on distortion estimation, characterized in that Including: Step S1: Obtain an original image, and extract image features from the original image to obtain an image feature map; Step S2: Input the image feature map into a distortion factor estimation module to estimate one or more factors causing the distortion of the original image, and obtain a distortion vector composed of the one or more factors causing the distortion of the original image; the distortion factor estimation module includes a convolutional layer, a global average pooling layer, and a fully connected layer, and each factor corresponds to a type of image distortion; Step S3: Input the distortion vector and the image feature map into an image reconstruction module, and the image reconstruction module removes the corresponding distortion and performs image reconstruction to obtain a residual image corresponding to the original image; based on the residual image, obtain a restored image corresponding to the original image; wherein, the image reconstruction module includes an attention convolutional block for removing distortion and performing image reconstruction; The image reconstruction module includes n attention convolutional blocks, a second convolutional layer, a third ReLU activation function layer, and a third convolutional layer connected in sequence; The attention convolutional block is used to remove distortion and perform image reconstruction; the attention convolutional block includes a reconstruction branch and an adjustment branch; The reconstruction branch is used to receive the image feature map or the feature map output by the previous-level attention convolutional block of the current attention convolutional block; if the current attention convolutional block is the first-level attention convolutional block, the input is the image feature map; if the current attention convolutional block is a non-first-level attention convolutional block, the input is the feature map output by the previous-level attention convolutional block of the current attention convolutional block; the reconstruction branch includes an attention sub-branch and a non-attention sub-branch, the attention sub-branch includes a fourth convolutional layer, a fourth ReLU activation function layer, a fifth convolutional layer, a spatial attention module, and a channel attention module connected in sequence, the attention sub-branch is used to remove distortion and reconstruct the image, and output an attention sub-branch residual feature map, the spatial attention module is used to enhance the image texture and detail information in the space and suppress the noise in the space, the channel attention module is used to enhance the channels containing the image texture and detail information and suppress the channels containing the interference information of the noise, and output an attention sub-branch residual feature map; wherein, the attention convolutional block is used to multiply with the output of the first distortion factor attention module of the adjustment branch to obtain a first fusion result; The non-attention sub-branch is used to receive the image feature map or the feature map output by the previous attention convolution block of the current attention convolution block; if the current attention convolution block is the first-level attention convolution block, the input is the image feature map; if the current attention convolution block is a non-first-level attention convolution block, the input is the feature map output by the previous attention convolution block of the current attention convolution block; the non-attention sub-branch includes a sixth convolutional layer, a fifth ReLU activation function layer, and a seventh convolutional layer connected in sequence. The non-attention sub-branch is used to supplement the feature information ignored by the attention sub-branch and output a non-attention sub-branch residual feature map. The non-attention sub-branch residual feature map is used to multiply the output of the second distortion factor attention module of the adjustment branch to obtain a second fusion result; The adjustment branch includes a first distortion factor attention module and a second distortion factor attention module. The input of the adjustment branch is the distortion vector. The first distortion factor attention module is used to generate a weight vector that controls the attention sub-branch of the reconstruction branch, and multiply the output with the output of the attention sub-branch to obtain a first fusion result; the second distortion factor attention module is used to generate a weight vector that controls the non-attention sub-branch of the reconstruction branch, and fuse the output with the output of the non-attention sub-branch to obtain a second fusion result; Add the first fusion result and the second fusion result to obtain a third fusion result, and add the third fusion result to the image feature map to obtain a processed image feature map, which is used as the output of the attention convolution block.

2. The method according to claim 1, wherein The distortion factor estimation module includes a first convolutional layer, a global average pooling layer, a first fully connected layer, a first ReLU activation function layer, a second fully connected layer, a second ReLU activation function layer, a third fully connected layer, and a Sigmoid activation function layer connected in sequence. Among them, the distortion factor estimation module inputs a distorted image. After being processed by the distortion factor estimation module, the distorted image outputs a distortion vector with the number of elements being k. The distortion factors are the respective elements in the vector, realizing the estimation of one or more factors that cause the distortion of the original image.

3. The method according to any one of claims 1-2, characterized in that, Based on the residual image, obtain the restored image corresponding to the original image, including: adding the obtained residual image and the original image pixel by pixel to obtain the restored image corresponding to the original image, and outputting the restored image.

4. An image restoration device based on distortion estimation, characterized in that The device includes: An image feature map acquisition module: configured to acquire an original image and extract image features from the original image to obtain an image feature map; A distortion estimation module: configured to input the image feature map into a distortion factor estimation module to estimate one or more factors that cause the distortion of the original image, and obtain a distortion vector composed of one or more factors that cause the distortion of the original image; the distortion factor estimation module includes a convolutional layer, a global average pooling layer, and a fully connected layer, and each factor corresponds to a type of image distortion; Image reconstruction module: Configured to input the distortion vector and the image feature map into the image reconstruction module. The image reconstruction module removes the corresponding distortion and performs image reconstruction to obtain the residual image corresponding to the original image. Based on the residual image, the restored image corresponding to the original image is obtained. Among them, the image reconstruction module includes an attention convolution block for removing distortion and performing image reconstruction. The image reconstruction module includes n attention convolution blocks, a second convolutional layer, a third ReLU activation function layer, and a third convolutional layer connected in sequence. The attention convolution block is used to remove distortion and perform image reconstruction. The attention convolution block includes a reconstruction branch and an adjustment branch. The reconstruction branch is used to receive the image feature map or the feature map output by the previous-level attention convolution block of the current attention convolution block. If the current attention convolution block is the first-level attention convolution block, the input is the image feature map. If the current attention convolution block is a non-first-level attention convolution block, the input is the feature map output by the previous-level attention convolution block of the current attention convolution block. The reconstruction branch includes an attention sub-branch and a non-attention sub-branch. The attention sub-branch includes a fourth convolutional layer, a fourth ReLU activation function layer, a fifth convolutional layer, a spatial attention module, and a channel attention module connected in sequence. The attention sub-branch is used to remove distortion and reconstruct the image, and outputs the attention sub-branch residual feature map. The spatial attention module is used to enhance the image texture and detail information in the space and suppress the noise in the space. The channel attention module is used to enhance the channels containing image texture and detail information and suppress the channels containing noise interference information, and outputs the attention sub-branch residual feature map. Among them, the attention convolution block is used to multiply with the output of the first distortion factor attention module of the adjustment branch to obtain the first fusion result. The non-attention sub-branch is used to receive the image feature map or the feature map output by the previous-level attention convolution block of the current attention convolution block. If the current attention convolution block is the first-level attention convolution block, the input is the image feature map. If the current attention convolution block is a non-first-level attention convolution block, the input is the feature map output by the previous-level attention convolution block of the current attention convolution block. The non-attention sub-branch includes a sixth convolutional layer, a fifth ReLU activation function layer, and a seventh convolutional layer connected in sequence. The non-attention sub-branch is used to supplement the feature information ignored by the attention sub-branch and outputs the non-attention sub-branch residual feature map. The non-attention sub-branch residual feature map is used to multiply with the output of the second distortion factor attention module of the adjustment branch to obtain the second fusion result. The adjustment branch includes a first distortion factor attention module and a second distortion factor attention module. The input of the adjustment branch is the distortion vector. The first distortion factor attention module is used to generate a weight vector for controlling the attention sub-branch of the reconstruction branch, multiply the output with the output of the attention sub-branch to obtain a first fusion result. The second distortion factor attention module is used to generate a weight vector for controlling the non-attention sub-branch of the reconstruction branch, and fuse the output with the output of the non-attention sub-branch to obtain a second fusion result. Add the first fusion result and the second fusion result to obtain a third fusion result, and add the third fusion result to the image feature map to obtain a processed image feature map, which is used as the output of the attention convolution block.

5. A computer-readable storage medium, in which multiple instructions are stored; the multiple instructions are used to be loaded and executed by a processor to perform the method according to any one of claims 1-3.

6. An electronic device, characterized in that, The electronic device includes: A processor for executing multiple instructions; A memory for storing multiple instructions; Wherein, the multiple instructions are used to be stored by the memory and loaded and executed by the processor to perform the method according to any one of claims 1-3.

Images