An all-in-one image restoration method based on deep learning

Abstract

The application provides an All-In-One image restoration method based on deep learning in the technical field of image restoration, comprising the following steps: S1, a large number of original images are acquired, each original image is subjected to randomization processing to obtain a corresponding noise image, and a data set is constructed based on each original image and noise image; S2, an image restoration model is created based on a prompt network, an encoder and a decoder, and a loss function of the image restoration model is set; 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; S3, the image restoration model is trained through the data set, and the hyperparameters of the image restoration model are continuously optimized in the training process; S4, the trained image restoration model is deployed, and the deployed image restoration model is used for image restoration. The application has the advantages that the generalization ability and the restoration quality of image restoration are greatly improved.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of image restoration, and particularly relates to an All-In-One image restoration method based on deep learning. BACKGROUND

[0002] During the acquisition process, images are affected by various unknown degradation types and degradation degrees, which reduces the image quality and even leads to the image being unusable, thus resulting in the need for image restoration. For image restoration, the traditional method can only restore clean images from images with a single degradation type or degradation degree, and cannot restore images from various unknown degradation types and degradation degrees, which leads to poor generalization ability of image restoration and affects the image restoration quality, thus failing to meet the current diversified 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 to be solved. SUMMARY

[0004] The present application solves the technical problem of providing an All-In-One image restoration method based on deep learning to improve the generalization ability and restoration quality of image restoration.

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

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

[0007] Step S2: creating an image restoration model based on a prompt network, an encoder and a decoder, 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; and the decoder is provided with a plurality of prompt blocks;

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

[0009] Step S4: deploying the trained image restoration model, and performing image restoration by using the deployed image restoration model.

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

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

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

[0013] Further, in the step S2, the data flow 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 is composed of a multi-head dual self-attention module and a gated feedforward neural network connected in sequence.

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

[0016] Further, the gated feedforward neural network is used to perform linear transformation on the input noise image through a first linear layer, divide the noise image into two image vectors of the same dimension through chunk operation, perform nonlinear activation on one of the image vectors through GELU activation function, and perform element multiplication on the other image vector and input a second linear layer to capture the association relationship between image features.

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

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

[0019] Further, the step S3 is specifically:

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

[0021] Further, the step S4 is specifically:

[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, a to-be-restored image is acquired, and the to-be-restored image is input into the image restoration model through the API interface to perform an image restoration operation.

[0023] The present application has the advantages that:

[0024] A large number of original images are acquired, randomization processing is performed on each original image 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 provided with a plurality of data flow architecture blocks, and a decoder provided with a plurality of prompt blocks, a 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 in the training process; finally, the trained image restoration model is deployed, and the deployed image restoration model is used for image restoration; that is, the image restoration model provided with the prompt network, the data flow architecture block, and the prompt block is used for image restoration, 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 double self-attention module and a gated feedforward neural network, the multi-head double self-attention module can better capture the dependency relationship between different spaces and channels in the noise 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 to improve the performance and learning ability of the model; based on the prompt information of the degradation type, the data flow architecture block encodes the degradation information implicitly and guides the network to distinguish different degradation types layer by layer, which is helpful for processing various unknown degradation; the loss function is constructed based on an L1 loss function, a cross-entropy loss function, and an aesthetic quality evaluation loss function, which further improves the perceptual quality and visual effect of the restored image; unlike the traditional pixel-based loss function, the loss function can better reflect the perception of the human eye to the image quality; in combination with the multi-scale network architecture of the image restoration model, the features of the noise image at different scales can be effectively captured and learned, which is helpful for processing image degradation of different sizes and types, and finally greatly improves the generalization ability and restoration quality of image restoration. BRIEF DESCRIPTION OF DRAWINGS

[0025] The application will be further described below with reference to the accompanying drawings and embodiments.

[0026] Figure 1 is a flowchart of an All-In-One image restoration method based on deep learning according to the present application.

[0027] Figure 2 is a schematic diagram of an image restoration model according to the present application.

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

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

[0030] Figure 5 is a schematic diagram of a gated feed-forward neural network (GFFN) of the present application.

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

[0032] The technical solutions in the embodiments of the present application have the following general idea: an image restoration model provided with a prompt network, a data stream architecture block, and a prompt block is used for image restoration, the prompt network is used for capturing prompt information of a degradation type from a noisy image, the data stream architecture block is composed of a multi-head dual self-attention module and a gated feed-forward neural network, the multi-head dual self-attention module can better capture the dependency relationship 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 feed-forward neural network can better utilize input features, thereby improving the performance and learning ability of the model, the data stream architecture block is based on the prompt information of the degradation type, and the degradation information is encoded implicitly to guide the network to distinguish different degradation types layer by layer, which is helpful for processing various unknown degradation, and the loss function is constructed based on an L1 loss function, a cross-entropy loss function, and an aesthetic quality evaluation loss function, thereby further improving the perceptual quality and visual effect of the restored image, which is different from the traditional pixel-based loss function and can better reflect the perception of the human eye to the quality of the image, and in combination with the multi-scale network architecture of the image restoration model, the features of the noisy image at different scales can be effectively captured and learned, which is helpful for processing 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 The preferred embodiment of an All-In-One image restoration method based on deep learning of the present application includes the following steps:

[0034] In step S1, a large number of original images are obtained, each original image is randomized to obtain a corresponding noisy image, and a dataset is constructed based on each original image and noisy image.

[0035] Step S2, create an image restoration model based on a prompt network, an encoder and a decoder, and set 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 (Dsa Blocks); and the decoder is provided with a plurality of prompt blocks (Prompt Blocks);

[0036] The high-resolution degraded image (noise image) obtains prompt information of the degradation type through the Prompt Net, increases the number of channels while gradually reducing the spatial resolution through a plurality of Dsa Blocks in the encoder, extracts low-resolution latent features, and then gradually recovers into a high-resolution clean output through the decoder; at each step of the decoder, the prompt block processed by the prompt information flow is embedded in the image to realize guided recovery.

[0037] Step S3, train the image restoration model through the data set, and constantly optimize the hyperparameters of the image restoration model during the training process;

[0038] Step S4, deploy the trained image restoration model, and use the deployed image restoration model for image restoration.

[0039] The application adopts a multi-scale architecture, uses a baseline model of Unet network, introduces a multi-scale feature fusion strategy, so that the image restoration model can process low-level detail information and high-level semantic information at the same time; this strategy enhances the learning ability of the image restoration model for different scale features in the noise image by integrating feature maps from different resolutions at different stages of the network.

[0040] The step S1 is specifically:

[0041] A large number of original images are obtained, and each original image is subjected to random processing of adding Gaussian noise, adding rain streaks or adding water mist to obtain a corresponding noise image, each noise image is labeled with a degradation type, and each original image and the corresponding noise image are associated to construct a data set based on each original image and noise image.

[0042] In step S2, the prompt network is used to capture prompt information of the degradation type from the noise image, and the prompt information is used to guide the model during the restoration process; the prompt network is a classification network based on Resnet50. The formula of the prompt network is:

[0043]

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

[0045] In the step S2, the data flow 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] To solve the problem that neural networks are difficult to capture the internal relationship of input data, the self-attention mechanism provides a solution. The core of the self-attention mechanism is to enhance the network's ability to perceive the relationship between different parts of the input data. Through self-attention, the network can identify and emphasize the mutual dependence between different regions in the input, thereby achieving more fine-grained feature extraction and information integration when processing complex data. That is, self-attention can be effectively applied between elements within a sequence to directly capture the dependency relationship within the sequence. Therefore, the present application uses a data flow architecture block based on the self-attention mechanism to deepen the understanding of the mutual dependence in the noise image and enhance the adaptability of the image restoration model to specific input content.

[0047] The data flow architecture block is composed of a multi-head double self-attention module (MHDSA) and a gated feed-forward neural network (GFFN) connected in sequence.

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

[0049] The present application uses MHDSA to replace the ordinary SA with quadratic complexity depending on the size of the input image resolution. MHDSA applies SA across space and channel dimensions through two parallel attention modules MHCSA and MHPSA, respectively, which can capture more rich context information, understand the global structure and local details of the image, understand complex image structures and long-distance dependencies, and improve overall performance.

[0050] MHCSA and MHPSA first use layer normalization (LN) to normalize each sample (noise image) individually, and then obtain the feature map U of the to-be-chunked through 1x1 convolution and 3x3 channel-wise 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 map is converted to the channel dimension M C ∈R C×N and the spatial dimension M P ∈R N×C , where N = H * W. The query (Q), key (K) and value (V) of the self-attention mechanism are obtained by chunking, and then the projection of Q and K is reshaped, and the attention matrix of different dimensions is obtained by calculating Q T *K, K is the dimension; the attention matrix is subjected to a softmax operation to obtain the relative factor that affects the dimension. After the softmax operation, the similarity matrix is multiplied by V, and through a 1x1 convolution and a residual connection with the original feature map, a feature representation that is re-encoded by the self-attention mechanism is obtained. The calculation formula of MHDSA is:

[0051]

[0052] where w1 and w2 represent the learnable parameter values of the fusion of the two feature maps; W 1×1 represents a 1x1 convolution; Attention C () represents the channel self-attention mechanism of different dimensions; Attention P () represents the spatial self-attention mechanism of different dimensions; alpha represents a learnable scaling parameter used to control and the size of the dot product before applying the softmax function; Y, X has the meaning given in Figure 4 ;

[0053] Similar to the traditional multi-head SA, the input feature map is divided into multiple "heads" in the channel and spatial dimensions, and separate attention maps are learned in parallel. For an input feature map U∈R H×W×C , the channel dimension is divided into h heads, and each "head" has C / h channels. Similarly, the spatial dimension is also divided.

[0054] The gating feedforward neural network is used to perform linear transformation on the input noise image through a first linear layer, and then the noise image is segmented into two image vectors with the same dimension through a chunk operation, one of the image vectors is subjected to non-linear activation through a GELU activation function, and then the other image vector is subjected to element multiplication 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 a complex feature transformation on the input representation through two layers of linear transformation and nonlinear activation, enhances the expression ability of the model, the first layer of linear transformation expands the input features, increases the dimension of the features, and thus captures more rich features and information, and the second layer of linear transformation maps the expanded features back to the original feature dimension.

[0056] The application enhances the learning ability of FFN representation through two important improvements of gating mechanism and channel-wise convolution. The gating mechanism divides the features after the first layer of linear transformation into two vectors with the same dimension through chunk operation, and realizes nonlinear activation through GELU activation function for one of the vectors, and inputs the element product of the two vectors into the second layer of linear transformation layer, and finely captures the complex interaction between features. In addition, 3*3 channel-wise convolution is also introduced in GFFN to ensure the independence of each channel feature, encode spatial dimension pixel position information, and realize higher quality image restoration. Through these improvements, GFFN can better utilize input features, improve the performance and learning ability of the model. Especially in processing complex image data, it shows significant advantages. The calculation formula of GFFN is:

[0057]

[0058] Wherein, φ() represents GELU activation function; LN() represents layer normalization; Represents element multiplication; The meaning of Figure 5 .

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

[0060] The image restoration model is used to restore the clean image from the unknown degraded image, although the image restoration model is initially unknown for the nature of degradation, but the performance of restoring the clean image can be enhanced by providing implicit context information about the degradation type. The application is based on prompt guided image restoration, and the prompt is an effective idea for processing All-In-One image restoration. While restoring the clean image, the related knowledge of the degradation type is used to increase the ability of the image restoration model to process different degradations. The application uses the idea of prompt, and proposes a component for All-In-One image restoration: prompt block.

[0061] PGM dynamically predicts attention-based weights from input features and applies them to P c To generate prompt weights from input features F, PGM first applies global average pooling in spatial dimensions to generate a feature vector, then reduces channels by a 1x1 convolutional layer to obtain a compact feature vector, followed by a softmax function to generate prompt weights, and finally uses prompt weights and prompt information to adjust P c and dynamically adjust through a 3x3 convolutional layer, the calculation formula of PGM is:

[0062]

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

[0064] where w i represents the prompt weight;

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

[0066]

[0067] where represents the connection operation.

[0068] In step S2, the loss function is constructed based on L1 loss function, cross-entropy loss function and aesthetic quality evaluation loss function to more comprehensively and effectively guide the optimization direction of 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] where ζ Rec (x,y) represents the L1 loss function; ζ CE (t,p) represents the cross-entropy loss function; ζ AQ (x,y) represents the aesthetic quality evaluation loss function.

[0072] The L1 loss function, also known as the absolute value loss function, 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 difference between the target value and the estimated value. Since the mean square error gives more weight to outliers, the model will try to reduce the error caused by outliers, thus 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] where x and y represent the input image and the true value image respectively; x i and y i represent 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] 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 features, thereby improving the effectiveness of the 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 represents the probability that Prompt Net predicts that sample i belongs to category C.

[0078] Image aesthetic quality evaluation aims to simulate human perception and evaluation of aesthetics, and takes the difference between the restored image and the true value as part of the loss function, which can improve the perceptual quality of image enhancement. The formula of the aesthetic quality evaluation 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:

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

[0083] The step S4 specifically includes:

[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, a to-be-restored image is acquired, and the to-be-restored image is input into the image restoration model through the API interface for image restoration operation.

[0085] In summary, the present application has the following advantages:

[0086] A large number of original images are acquired, and corresponding noise images are obtained by randomizing each original image, and a data set is constructed based on each original image and noise image; then an image restoration model is created based on a prompt network, an encoder provided with a plurality of data flow architecture blocks, and a decoder provided with a plurality of prompt blocks, and a 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 for image restoration; that is, the image restoration model provided with the prompt network, the data flow architecture block and the prompt block is used for image restoration, 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 double self-attention module and a gated feedforward neural network, the multi-head double self-attention module can better capture the dependency between different spaces and channels in the noise 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 to improve the performance and learning ability of the model; the data flow architecture block encodes the degradation information implicitly based on the prompt information of the degradation type, and guides the network to distinguish different degradation types layer by layer, which is helpful for processing various unknown degradation; the loss function is constructed based on the L1 loss function, the cross-entropy loss function and the aesthetic quality evaluation 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 perception of human eyes to image quality; combined with the multi-scale network architecture of the image restoration model, the features of the noise image at different scales can be effectively captured and learned, which is helpful for processing image degradation of different sizes and types, and finally greatly improves the generalization ability and restoration quality of image restoration.

[0087] While the foregoing description has described specific embodiments of the application, one ordinary skill in the art will appreciate that various modifications and changes can be made thereto without departing from the spirit and scope of the application, as set forth in the appended claims.

Claims

1. A deep learning-based all-in-one image restoration method, characterized in that: Includes the following steps: Step S1: Obtain a large number of original images, randomize each of the original images to obtain corresponding noise images, and construct a dataset based on each of the original images and the noise images; Step S2: Create an image restoration model based on the cue network, encoder, and decoder, and set the loss function of the image restoration model; the cue network, encoder, and decoder are connected in sequence; the encoder has several data flow architecture blocks; the decoder has several cue blocks; Step S3: Train the image restoration model using the dataset, and continuously optimize the hyperparameters of the image restoration model during the training process; Step S4: Deploy the trained image restoration model and use the deployed image restoration model to perform image restoration; In step S2, the cue network is used to capture degraded cue information from the noisy image and is built based on ResNet50; the dataflow architecture block is used to gradually reduce the spatial resolution of the noisy image while increasing the number of channels in order to extract latent features. The prompt block consists of a prompt generation module and a prompt interaction module; the loss function is constructed based on the L1 loss function, the cross-entropy loss function, and the aesthetic quality assessment loss function. The dataflow architecture block consists of a multi-head dual self-attention module and a gated feedforward neural network connected in sequence. The multi-head dual self-attention module consists of a multi-head attention unit MHCSA for cross-space application of the spatial attention module and a multi-head attention unit MHPSA for cross-channel application of the spatial attention module. The gated feedforward neural network is used to perform a linear transformation on the input noisy image through a first linear layer, and then segment the noisy image into two image vectors with the same dimension through a chunk operation. One of the image vectors is non-linearly activated by the 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.

2. The all-in-one image restoration method based on deep learning as described in claim 1, characterized in that: Step S1 specifically involves: A large number of original images are acquired, and each original image is randomized by adding Gaussian noise, rain streaks, or water mist to obtain a corresponding noise image. Each noise image is labeled with a degradation type. After associating each original image with the corresponding noise image, a dataset is constructed based on each original image and the noise image.

3. The all-in-one image restoration method based on deep learning as described in claim 1, characterized in that: Step S3 specifically involves: The dataset is divided into a training set and a validation set based on a preset ratio. 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. The trained image restoration model is then validated using the validation set.

4. The all-in-one image restoration method based on deep learning as described in claim 1, characterized in that: Step S4 specifically involves: The trained image restoration model is deployed on a cloud server. The API interface of the image restoration model deployed on the cloud server is set to call the image restoration model to obtain the image to be restored. The image to be restored is input into the image restoration model through the API interface to perform the image restoration operation.

Images