All-In-One image restoration method based on deep learning

Abstract

The invention provides an All-In-One image restoration method based on deep learning in the technical field of image restoration, and the method comprises the steps: S1, obtaining a large number of original images, carrying out the randomization processing of each original image, obtaining a corresponding noise image, and constructing a data set based on each original image and the noise image; s2, creating an image restoration model based on the prompt network, the encoder and the decoder, and setting a loss function of the image restoration model; the prompt network, the encoder and the decoder are connected in sequence; the encoder is provided with a plurality of data stream architecture blocks; the decoder is provided with a plurality of prompt blocks; s3, training the image restoration model through the data set, and continuously optimizing hyper-parameters of the image restoration model in the training process; and S4, deploying the trained image restoration model, and performing image restoration by using the deployed image restoration model. The method has the advantages that the generalization ability and the restoration quality of image restoration are greatly improved.

Description

Technical Field

[0001] The present invention relates to the technical field of image restoration, and in particular to an All-In-One image restoration method based on deep learning. Background Art

[0002] During the acquisition process, images are affected by a variety of unknown degradation types and degrees, which will reduce the image quality or even make the image unusable, thus creating a need for image restoration. Traditionally, image restoration can only restore clean images from images of a single degradation type or degree, and cannot restore images from various unknown degradation types and degrees. This results in poor generalization of image restoration, which in turn affects the quality of image restoration and cannot meet the current diverse image restoration needs.

[0003] Therefore, how to provide an All-In-One image restoration method based on deep learning to improve the generalization ability and restoration quality of image restoration has become a technical problem that needs to be solved urgently. Summary of the invention

[0004] The technical problem to be solved by the present invention is to provide an All-In-One image restoration method based on deep learning, so as to improve the generalization ability and restoration quality of image restoration.

[0005] The present invention is implemented as follows: an All-In-One image restoration method based on deep learning comprises the following steps:

[0006] Step S1, obtaining a large number of original images, performing random processing on each of the original images to obtain a corresponding noise image, and constructing a data set based on each of the original images and the noise image;

[0007] Step S2, creating an image restoration model based on the prompt network, the encoder and the decoder, and setting the loss function of the image restoration model; the prompt network, the encoder and the decoder are connected in sequence; the encoder is provided with a plurality of data flow architecture blocks; the decoder is provided with a plurality of prompt blocks;

[0008] Step S3, training the image restoration model using the data set, and continuously optimizing the hyperparameters of the image restoration model during the training process;

[0009] Step S4: deploy the trained image restoration model, and perform image restoration using the deployed image restoration model.

[0010] Furthermore, the step S1 is specifically as follows:

[0011] A large number of original images are obtained, and Gaussian noise, rain streaks or water mist are added to each of the original images to obtain a corresponding noise image. The degradation type of each noise image is labeled, and after associating each of the original images with the corresponding noise image, a data set is constructed based on each of the original images and the noise image.

[0012] Furthermore, in step S2, the prompt network is used to capture the prompt information of the degradation type from the noisy image and is constructed based on Resnet50.

[0013] Furthermore, in step S2, the data stream architecture block is used to increase the number of channels while gradually reducing the spatial resolution of the noise image to extract potential features;

[0014] The data flow architecture block consists of a multi-head dual self-attention module and a gated feedforward neural network connected in sequence.

[0015] Furthermore, the multi-head dual self-attention module is composed of a multi-head attention unit MHCSA for applying the spatial attention module across space and a multi-head attention unit MHPSA for applying the spatial attention module across channels.

[0016] Furthermore, the gated feedforward neural network is used to perform a linear transformation on the input noise image through a first linear layer, and then divide the noise image into two image vectors of the same dimension through a chunk operation, wherein one of the image vectors is nonlinearly activated by a GELU activation function, and then element-wise multiplied with the other image vector and input into a second linear layer to capture the correlation between image features.

[0017] Furthermore, in step S2, the prompt block is composed of a prompt generation module and a prompt interaction module.

[0018] Furthermore, in step S2, the loss function is constructed based on the L1 loss function, the cross entropy loss function and the aesthetic quality assessment loss function.

[0019] Furthermore, the step S3 is specifically as follows:

[0020] The data set is divided into a training set and a validation set based on a preset ratio, and the image restoration model is trained using the training set until the loss value of the loss function is less than a preset loss threshold. During the training process, the hyperparameters of the image restoration model are continuously optimized, and the trained image restoration model is verified using the validation set.

[0021] Furthermore, the step S4 is specifically as follows:

[0022] The trained image restoration model is deployed on a cloud server, an API interface for calling the image restoration model deployed on the cloud server is set, an image to be restored is obtained, and the image to be restored is input into the image restoration model through the API interface to perform an image restoration operation.

[0023] The advantages of the present invention are:

[0024] A large number of original images are obtained, each original image is randomly processed to obtain a corresponding noise image, and a data set is constructed based on each original image and the noise image; then an image restoration model is created based on a prompt network, an encoder with several data flow architecture blocks, and a decoder with several prompt blocks, and the loss function of the image restoration model is set; the image restoration model is trained through the data set, and the hyperparameters of the image restoration model are continuously optimized during the training process; finally, the trained image restoration model is deployed, and the deployed image restoration model is used to perform image restoration; that is, image restoration is performed through an image restoration model with a prompt network, a data flow architecture block, and a prompt block, and the prompt network is used to capture the prompt information of the degradation type from the noise image; the data flow architecture block is composed of a multi-head dual self-attention module and a gated feedforward neural network, and the multi-head dual self-attention module uses multi-head spatial self-attention and multi-head channel Self-attention can better capture the dependencies between different spaces and channels in noisy images, thereby improving the quality of feature learning. The gated feedforward neural network can better utilize the input features and improve the performance and learning ability of the model. The data flow architecture block is based on the prompt information of the degradation type. By implicitly encoding the degradation information, the network is guided layer by layer to distinguish different degradation types, which is helpful to deal with a variety of unknown degradations. The loss function is constructed based on the L1 loss function, the cross entropy loss function and the aesthetic quality assessment loss function to further improve the perceptual quality and visual effect of the restored image. Unlike the traditional pixel-based loss function, it can better reflect the human eye's perception of image quality. Combined with the multi-scale network architecture of the image restoration model, it can effectively capture and learn the features of noisy images at different scales, which is helpful to deal with image degradation of different sizes and types, and ultimately greatly improves the generalization ability and restoration quality of image restoration. BRIEF DESCRIPTION OF THE DRAWINGS

[0025] The present invention will be further described below in conjunction with embodiments with reference to the accompanying drawings.

[0026] Figure 1 It is a flow chart of an All-In-One image restoration method based on deep learning of the present invention.

[0027] Figure 2 Schematic diagram of the image restoration model of the present invention.

[0028] Figure 3It is a schematic diagram of the data flow architecture block (Dsa Block) of the present invention.

[0029] Figure 4 It is a schematic diagram of the multi-head dual self-attention module (MHDSA) of the present invention.

[0030] Figure 5 Schematic diagram of the gated feedforward neural network (GFFN) of the present invention.

[0031] Figure 6 It is a schematic diagram of the prompt block (Prompt Block) of the present invention. DETAILED DESCRIPTION

[0032] The technical solution in the embodiment of the present application has the following general idea: image restoration is performed by an image restoration model provided with a prompt network, a data flow architecture block and a prompt block, wherein the prompt network is used to capture the prompt information of the degradation type from the noisy image; the data flow architecture block is composed of a multi-head dual self-attention module and a gated feedforward neural network, and the multi-head dual self-attention module can better capture the dependencies between different spaces and channels in the noisy image through multi-head spatial self-attention and multi-head channel self-attention, thereby improving the quality of feature learning, and the gated feedforward neural network can better utilize the input features and improve the performance and learning ability of the model; the data flow architecture block is based on The hint information of the degradation type guides the network to distinguish different degradation types layer by layer by implicitly encoding the degradation information, which helps to deal with various unknown degradations; the loss function is constructed based on the L1 loss function, the cross entropy loss function and the aesthetic quality assessment loss function, which further improves the perceptual quality and visual effect of the restored image. Unlike the traditional pixel-based loss function, it can better reflect the human eye's perception of image quality; combined with the multi-scale network architecture of the image restoration model, it can effectively capture and learn the features of noisy images at different scales, which helps to deal with image degradation of different sizes and types, thereby improving the generalization ability and restoration quality of image restoration.

[0033] Please refer to Figures 1 to 6 As shown, a preferred embodiment of the All-In-One image restoration method based on deep learning of the present invention comprises the following steps:

[0034] Step S1, obtaining a large number of original images, performing random processing on each of the original images to obtain a corresponding noise image, and constructing a data set based on each of the original images and the noise image;

[0035] Step S2, creating an image restoration model based on a prompt network (Prompt Net), an encoder and a decoder, and setting a loss function of the image restoration model; the prompt network, the encoder and the decoder are connected in sequence; the encoder is provided with a plurality of data flow architecture blocks (Dsa Block); the decoder is provided with a plurality of prompt blocks (Prompt Block);

[0036] The high-resolution degraded image (noise image) obtains the prompt information of the degradation type through Prompt Net, and gradually reduces the spatial resolution while increasing the number of channels through multiple Dsa Blocks in the encoder to extract low-resolution potential features, and then gradually restores it to high-resolution clean output through the decoder; at each step of the decoder, the image is embedded with a prompt block processed by the prompt information stream to achieve guided recovery.

[0037] Step S3, training the image restoration model using the data set, and continuously optimizing the hyperparameters of the image restoration model during the training process;

[0038] Step S4: deploy the trained image restoration model, and perform image restoration using the deployed image restoration model.

[0039] The present invention adopts a multi-scale architecture and uses the Unet network as the baseline model. By introducing a multi-scale feature fusion strategy, the image restoration model can simultaneously process low-level detail information and high-level semantic information; this strategy enhances the image restoration model's ability to learn features of different scales in noisy images by integrating feature maps from different resolutions at different stages of the network.

[0040] The step S1 is specifically as follows:

[0041] A large number of original images are obtained, and Gaussian noise, rain streaks or water mist are added to each of the original images to obtain a corresponding noise image. The degradation type of each noise image is labeled, and after associating each of the original images with the corresponding noise image, a data set is constructed based on each of the original images and the noise image.

[0042] In step S2, the hint network is used to capture hint information of degradation type from the noisy image, and use the hint information to guide the model during the restoration process. It is an image classification network built based on Resnet50. The formula of the hint network is:

[0043]

[0044] Where F represents the input feature; represents the output feature; P crepresents the learnable parameters Prompt components, i.e., a set of learnable parameters that interact with the incoming features and prompt information to embed degradation information; P i represents prompt information; PGM() represents prompt generation module (PGM); PIM() represents prompt interaction module (PIM);

[0045] In step S2, the data stream architecture block is used to increase the number of channels while gradually reducing the spatial resolution of the noise image to extract potential features;

[0046] The self-attention mechanism provides a solution to the problem that neural networks have difficulty capturing the internal relationship of input data. The core of the self-attention mechanism is to enhance the network's perception of the relationship between the various parts of the input data. Through self-attention, the network can identify and emphasize the interdependence between different regions in the input, thereby achieving more sophisticated feature extraction and information integration when processing complex data. That is, self-attention can be effectively applied to the elements within the sequence, directly capturing the dependencies within the sequence. Therefore, the present invention uses a data flow architecture block based on the self-attention mechanism to deepen the understanding of interdependence in noisy images and enhance the adaptability of the image restoration model to specific input content.

[0047] The data flow architecture block consists of a multi-head dual self-attention module (MHDSA) and a gated feedforward neural network (GFFN) connected in sequence.

[0048] The multi-head dual self-attention module is composed of a multi-head attention unit MHCSA for applying the spatial attention module across space and a multi-head attention unit MHPSA for applying the spatial attention module across channels.

[0049] The present invention uses MHDSA to replace the ordinary SA which has quadratic complexity with the input image resolution. MHDSA uses SA across spatial and channel dimensions through two parallel attention modules MHCSA and MHPSA, respectively, which can capture richer contextual information, understand the global structure and local details of the image, understand the complex image structure and long-distance dependencies, and improve the overall performance.

[0050] MHCSA and MHPSA first use layer normalization (LN) to normalize each sample (noise image) separately, and then obtain the feature map U∈R to be divided into chunks through 1×1 convolution and 3×3 channel-by-channel convolution. H×W×C, R is a real number, H is the image height, W is the image width, and C is the number of channels; according to the different dimensions of the two self-attention mechanisms, the feature maps are converted to the channel dimension M respectively. C ∈R C×N and the spatial dimension M P ∈R N×C , where N = H*W. The query (Query, Q), key (Key, K) and value (Value, V) of the self-attention mechanism are obtained through chunks. Next, the projections of Q and K are reshaped by calculating Q T *K obtains attention matrices of different dimensions, where K is the dimension; performs a softmax operation on the attention matrix to obtain the relative factors that affect each dimension. The similarity matrix after the softmax operation is multiplied by V, and connected through a 1×1 convolution and the residual of the original feature map to obtain the feature representation re-encoded by the self-attention mechanism. The calculation formula of MHDSA is:

[0051]

[0052] Among them, w1 and w2 represent the learnable parameter values ​​of the fusion of two feature maps respectively; W 1×1 Represents 1×1 convolution; Attention C () represents the channel self-attention mechanism of different dimensions; Attention P () represents the spatial self-attention mechanism in different dimensions; α represents a learnable scaling parameter used to control the and The size of the dot product; Y, The meaning of X is shown in Figure 4 ;

[0053] Similar to the traditional multi-head SA, this paper divides the input feature map into multiple “heads” in the channel and spatial dimensions, and learns separate attention maps in parallel. H×W×C , the present invention divides its channel dimension into h heads, each "head" has C / h channels. Similarly, the division is also performed in the spatial dimension.

[0054] The gated feedforward neural network is used to perform a linear transformation on the input noise image through a first linear layer, and then divide the noise image into two image vectors of the same dimension through a chunk operation, wherein one of the image vectors is nonlinearly activated by a GELU activation function, and then element-wise multiplied with the other image vector and input into a second linear layer to capture the correlation between image features.

[0055] In order to transform features, the traditional feed-forward network (FFN) receives the output from the attention mechanism and performs complex feature transformation on the input representation through two layers of linear transformation and nonlinear activation to enhance the expressiveness of the model. The first layer of linear transformation expands the input features and increases the dimension of the features to capture richer features and information. The second layer of linear transformation maps the expanded features back to the original feature dimensions.

[0056] The present invention enhances the learning ability of FFN representation through two important improvements: gating mechanism and channel-by-channel convolution. The gating mechanism uses chunk operation to divide the features after the first layer of linear transformation into two vectors of the same dimension, and implements nonlinear activation of one of the vectors through the GELU activation function. The two vectors are element-wise multiplied and input into the second layer of linear transformation layer to capture the complex interactions between features in detail. In addition, 3×3 channel-by-channel convolution is also introduced in GFFN to ensure the independence of each channel feature, while encoding the pixel position information in the spatial dimension to achieve higher quality image restoration. Through these improvements, GFFN can better utilize the input features and improve the performance and learning ability of the model. In particular, it shows significant advantages when processing complex image data. The calculation formula of GFFN is:

[0057]

[0058] Among them, φ() represents the GELU activation function; LN() represents layer normalization; represents element-wise multiplication; See the meaning of Figure 5 .

[0059] In the step S2, the prompt block is composed of a prompt generation module (Prompt Generation Module, PGM) and a prompt interaction module (Prompt Interaction Module, PIM).

[0060] The image restoration model is used to restore a clean image from an unknown degraded image. Although the image restoration model is initially unknown to the nature of the degradation, its performance in restoring the clean image can be enhanced by providing implicit contextual information about the degradation type. The present invention is based on prompt-guided image restoration. Prompts are an idea for effectively processing All-In-One image restoration. While restoring the clean image, the relevant knowledge of the degradation type is used to increase the ability of the image restoration model to handle different degradations. The present invention uses the idea of ​​prompts to propose a component for All-In-One image restoration: a prompt block.

[0061] PGM dynamically predicts attention-based weights from the input features and applies them to P c To generate the input prompt P, in order to generate the prompt weight from the input feature F, PGM first applies global average pooling in the spatial dimension to generate the feature vector, then reduces the channel through a 1×1 convolutional layer to obtain a compact feature vector, and then generates the prompt weight through the softmax function, and finally uses the prompt weight and the prompt information to P c Adjust and dynamically adjust through 3×3 convolutional layers. The calculation formula of PGM is:

[0062]

[0063] w i =softmax(Conv 1×1 (GAP(F)));

[0064] Among them, w i represents the prompt weight;

[0065] The main goal of PIM is to build the association between the input features F and the input prompts P to achieve guided recovery. In PIM, the generated input prompts P are concatenated with the input features F along the channel dimension, and then the representation is passed through the Dsa Block to transform the input features using the degradation information encoded in the input prompts P. The calculation formula of PIM is:

[0066]

[0067] Among them, ; represents the connection operation.

[0068] In step S2, the loss function is constructed based on the L1 loss function, the cross entropy loss function and the aesthetic quality assessment loss function to more comprehensively and effectively guide the optimization direction of the network parameters.

[0069] The formula of the loss function is:

[0070] ζ total (x,y,t,p)=ζ Rec (x,y)+ζ CE (t,p)+ζ AQ (x,y);

[0071] Among them, Rec (x,y) represents the L1 loss function; ζ CE (t,p) represents the cross entropy loss function; ζ AQ (x,y) represents the aesthetic quality assessment loss function;

[0072] The L1 loss function is also called the absolute value loss function. It is used to calculate the pixel-level content reconstruction loss between the restored image and the clean image. The optimization goal is to minimize the sum of the absolute value differences between the target value and the estimated value. Since the mean square error gives greater weight to outliers, the model will strive to reduce the error caused by outliers, thereby reducing the overall performance of the model. Therefore, when training the network, the absolute value loss is more effective for image restoration tasks. The formula of the L1 loss function is:

[0073]

[0074] Among them, x and y represent the input image and the true value image respectively; x i and i represents the i-th input image and the true value image; T represents the number of pixels in the image; i represents the index of the pixel in the image; f() represents the restoration network.

[0075] In order to make the prompt information output by Prompt Net more accurate and effective, a cross entropy loss function is added to further optimize the prompt generation process. By minimizing the cross entropy loss, the generated prompt information can more accurately reflect the degradation category of the input feature, thereby improving the effectiveness of prompt information encoding. The formula of the cross entropy loss function is:

[0076]

[0077] Where C represents the total number of categories; t i,c represents the true label of sample i. If sample i belongs to category C, then t i,c =1, otherwise t i,c =0; p i,c It represents the probability that Prompt Net predicts that sample i belongs to category C.

[0078] Image aesthetic quality assessment aims to simulate human perception and evaluation of aesthetics. Taking the difference between the restored image and the true value as part of the loss function can improve the perceptual quality of image enhancement. The formula of the aesthetic quality assessment loss function is:

[0079]

[0080] Where K represents the range of aesthetic score; p k () represents the probability of each rating.

[0081] The step S3 is specifically as follows:

[0082] The data set is divided into a training set and a validation set based on a preset ratio, and the image restoration model is trained using the training set until the loss value of the loss function is less than a preset loss threshold. During the training process, the hyperparameters of the image restoration model are continuously optimized, and the trained image restoration model is verified using the validation set.

[0083] The step S4 is specifically as follows:

[0084] The trained image restoration model is deployed on a cloud server, an API interface for calling the image restoration model deployed on the cloud server is set, an image to be restored is obtained, and the image to be restored is input into the image restoration model through the API interface to perform an image restoration operation.

[0085] In summary, the advantages of the present invention are:

[0086] A large number of original images are obtained, each original image is randomly processed to obtain a corresponding noise image, and a data set is constructed based on each original image and the noise image; then an image restoration model is created based on a prompt network, an encoder with several data flow architecture blocks, and a decoder with several prompt blocks, and the loss function of the image restoration model is set; the image restoration model is trained through the data set, and the hyperparameters of the image restoration model are continuously optimized during the training process; finally, the trained image restoration model is deployed, and the deployed image restoration model is used to perform image restoration; that is, image restoration is performed through an image restoration model with a prompt network, a data flow architecture block, and a prompt block, and the prompt network is used to capture the prompt information of the degradation type from the noise image; the data flow architecture block is composed of a multi-head dual self-attention module and a gated feedforward neural network, and the multi-head dual self-attention module uses multi-head spatial self-attention and multi-head channel Self-attention can better capture the dependencies between different spaces and channels in noisy images, thereby improving the quality of feature learning. The gated feedforward neural network can better utilize the input features and improve the performance and learning ability of the model. The data flow architecture block is based on the prompt information of the degradation type. By implicitly encoding the degradation information, the network is guided layer by layer to distinguish different degradation types, which is helpful to deal with a variety of unknown degradations. The loss function is constructed based on the L1 loss function, the cross entropy loss function and the aesthetic quality assessment loss function to further improve the perceptual quality and visual effect of the restored image. Unlike the traditional pixel-based loss function, it can better reflect the human eye's perception of image quality. Combined with the multi-scale network architecture of the image restoration model, it can effectively capture and learn the features of noisy images at different scales, which is helpful to deal with image degradation of different sizes and types, and ultimately greatly improves the generalization ability and restoration quality of image restoration.

[0087] Although the specific implementation modes of the present invention are described above, those skilled in the art should understand that the specific implementation modes described are only illustrative and are not intended to limit the scope of the present invention. Equivalent modifications and changes made by those skilled in the art in accordance with the spirit of the present invention should be included in the scope of protection of the claims of the present invention.

Claims

1. An All-In-One image restoration method based on deep learning, characterized in that: The steps include: Step S1, obtaining a large number of original images, performing random processing on each of the original images to obtain a corresponding noise image, and constructing a data set based on each of the original images and the noise image; Step S2, creating an image restoration model based on the prompt network, the encoder and the decoder, and setting the loss function of the image restoration model; the prompt network, the encoder and the decoder are connected in sequence; the encoder is provided with a plurality of data flow architecture blocks; the decoder is provided with a plurality of prompt blocks; Step S3, training the image restoration model using the data set, and continuously optimizing the hyperparameters of the image restoration model during the training process; Step S4: deploy the trained image restoration model, and perform image restoration using the deployed image restoration model.

2. The All-In-One image restoration method based on deep learning as claimed in claim 1, characterized in that: The step S1 is specifically as follows: A large number of original images are obtained, and Gaussian noise, rain streaks or water mist are added to each of the original images to obtain a corresponding noise image. The degradation type of each noise image is labeled, and after associating each of the original images with the corresponding noise image, a data set is constructed based on each of the original images and the noise image.

3. The All-In-One image restoration method based on deep learning as claimed in claim 1, characterized in that: In step S2, the prompt network is used to capture the prompt information of the degradation type from the noisy image and is constructed based on Resnet50.

4. The All-In-One image restoration method based on deep learning as claimed in claim 1, characterized in that: In step S2, the data stream architecture block is used to increase the number of channels while gradually reducing the spatial resolution of the noise image to extract potential features; The data flow architecture block consists of a multi-head dual self-attention module and a gated feedforward neural network connected in sequence.

5. The All-In-One image restoration method based on deep learning as claimed in claim 4, characterized in that: The multi-head dual self-attention module is composed of a multi-head attention unit MHCSA for applying the spatial attention module across space and a multi-head attention unit MHPSA for applying the spatial attention module across channels.

6. The All-In-One image restoration method based on deep learning as claimed in claim 4, characterized in that: The gated feedforward neural network is used to perform a linear transformation on the input noise image through a first linear layer, and then divide the noise image into two image vectors of the same dimension through a chunk operation, wherein one of the image vectors is nonlinearly activated by a GELU activation function, and then element-wise multiplied with the other image vector and input into a second linear layer to capture the correlation between image features.

7. The All-In-One image restoration method based on deep learning as claimed in claim 1, characterized in that: In the step S2, the prompt block is composed of a prompt generation module and a prompt interaction module.

8. The All-In-One image restoration method based on deep learning as claimed in claim 1, characterized in that: In step S2, the loss function is constructed based on the L1 loss function, the cross entropy loss function and the aesthetic quality assessment loss function.

9. The All-In-One image restoration method based on deep learning as claimed in claim 1, characterized in that: The step S3 is specifically as follows: The data set is divided into a training set and a validation set based on a preset ratio, and the image restoration model is trained using the training set until the loss value of the loss function is less than a preset loss threshold. During the training process, the hyperparameters of the image restoration model are continuously optimized, and the trained image restoration model is verified using the validation set.

10. The All-In-One image restoration method based on deep learning according to claim 1, characterized in that: The step S4 is specifically as follows: The trained image restoration model is deployed on a cloud server, an API interface for calling the image restoration model deployed on the cloud server is set, an image to be restored is obtained, and the image to be restored is input into the image restoration model through the API interface to perform an image restoration operation.

Images