Training method of image restoration model and image restoration method and device

Abstract

The invention relates to the technical field of image processing, and discloses a training method of an image restoration model, and an image restoration method and device, and the method comprises the steps: determining a first image feature corresponding to a first image based on the image restoration model, and carrying out the optimization of a plurality of preset discrete vectors in a discrete codebook based on the first image feature, obtaining a target discrete codebook; determining a second image feature corresponding to the second image based on the image restoration model, and determining a plurality of prediction discrete vectors corresponding to the second image feature in the target discrete codebook based on a prediction module in the image restoration model; the target discrete vector comprises a prediction discrete vector; and optimizing the prediction module based on the plurality of target discrete vectors and the plurality of prediction discrete vectors to obtain a target prediction module, and determining a trained target image restoration model based on the target discrete codebook and the target prediction module. According to the invention, the quality of image restoration can be improved, and the robustness of image restoration in an image degradation scene can be enhanced.

Description

Technical Field

[0001] The present application relates to the field of image processing technology, and in particular to a training method for an image restoration model, an image restoration method and a device. Background Art

[0002] At present, images captured in the real world are often affected by various image degradation phenomena, such as compression artifacts, blur, noise, etc. These image degradation phenomena seriously affect the image quality, resulting in the degradation of the performance of image recognition, analysis, processing and other technologies. Therefore, researchers have developed image restoration (or image restoration) technology for degraded images, aiming to restore degraded low-quality images to generate high-quality images.

[0003] However, existing image restoration methods usually lack the ability to restore image details and are difficult to maintain high image fidelity. Therefore, existing image restoration methods still need further improvement. Summary of the invention

[0004] The embodiments of the present application provide a training method for an image restoration model, an image restoration method and a device, which can improve the quality of image restoration and enhance the robustness and generalization ability of image restoration in image degradation scenarios. Specifically, the embodiments of the present application disclose the following technical solutions:

[0005] A first aspect of an embodiment of the present application provides a training method for an image restoration model, the method comprising: determining a first image feature corresponding to a first image based on an image restoration model, and optimizing multiple preset discrete vectors in a discrete codebook based on the first image feature to obtain a target discrete codebook; wherein the target discrete codebook includes multiple target discrete vectors; determining a second image feature corresponding to a second image based on the image restoration model, and determining multiple predicted discrete vectors corresponding to the second image feature in the target discrete codebook based on a prediction module in the image restoration model; wherein a quality parameter of the first image is higher than a quality parameter of the second image, and the target discrete vector includes the predicted discrete vector; optimizing the prediction module based on the multiple target discrete vectors and the multiple predicted discrete vectors to obtain a target prediction module, and determining a trained target image restoration model based on the target discrete codebook and the target prediction module.

[0006] In some embodiments, the above-mentioned optimizing multiple preset discrete vectors in the discrete codebook based on the above-mentioned first image feature to obtain the target discrete codebook includes: determining multiple matching discrete vectors corresponding to the above-mentioned first image feature in the above-mentioned discrete codebook; wherein the above-mentioned preset discrete vectors include the above-mentioned matching discrete vectors; based on the above-mentioned first image feature and the above-mentioned multiple matching discrete vectors, optimizing the above-mentioned multiple preset discrete vectors in the above-mentioned discrete codebook to obtain the above-mentioned target discrete codebook.

[0007] In some embodiments, the above-mentioned determination of the first image feature corresponding to the first image based on the image restoration model includes: encoding the above-mentioned first image based on the first encoder in the above-mentioned image restoration model to obtain the above-mentioned first image feature and multiple first image sub-features corresponding to the above-mentioned first image feature; the above-mentioned determination of the multiple matching discrete vectors corresponding to the above-mentioned first image feature in the above-mentioned discrete codebook includes: determining the candidate similarities between the above-mentioned first image sub-feature and each of the above-mentioned preset discrete vectors in the above-mentioned multiple preset discrete vectors; determining the target similarity with the highest similarity among the multiple candidate similarities, and determining the target preset discrete vector corresponding to the above-mentioned target similarity as the above-mentioned matching discrete vector corresponding to the above-mentioned first image sub-feature; wherein the above-mentioned preset discrete vector includes the above-mentioned target preset discrete vector.

[0008] In some embodiments, the above-mentioned determination of multiple predicted discrete vectors corresponding to the above-mentioned second image features in the above-mentioned target discrete codebook based on the prediction module in the above-mentioned image restoration model includes: determining multiple prediction vector indexes corresponding to the above-mentioned second image features in the above-mentioned target discrete codebook based on the above-mentioned prediction module; wherein the above-mentioned target discrete codebook includes target vector indexes corresponding to each of the above-mentioned target discrete vectors, and the above-mentioned target vector indexes include the above-mentioned prediction vector indexes; determining the above-mentioned multiple prediction discrete vectors corresponding to the above-mentioned second image features in the above-mentioned target discrete codebook based on the above-mentioned multiple prediction vector indexes.

[0009] In some embodiments, the above-mentioned determining the multiple prediction vector indexes corresponding to the above-mentioned second image features in the above-mentioned target discrete codebook based on the above-mentioned prediction module includes: determining the high-level features corresponding to the above-mentioned second image features based on the above-mentioned prediction module; wherein the above-mentioned high-level features are determined using a self-attention mechanism; and determining the above-mentioned multiple prediction vector indexes corresponding to the above-mentioned high-level features in the above-mentioned target discrete codebook based on the above-mentioned prediction module.

[0010] In some embodiments, the above-mentioned prediction module is optimized based on the above-mentioned multiple target discrete vectors and the above-mentioned multiple prediction discrete vectors to obtain the target prediction module, including: determining a code-level loss function based on the above-mentioned multiple target vector indices corresponding to the above-mentioned multiple target discrete vectors and the above-mentioned multiple prediction vector indices corresponding to the above-mentioned multiple prediction discrete vectors; optimizing the above-mentioned prediction module based on the above-mentioned code-level loss function to obtain the above-mentioned target prediction module.

[0011] In some embodiments, determining the second image features corresponding to the second image based on the image restoration model includes: encoding the second image based on the second encoder in the image restoration model to obtain the second image features corresponding to the second image.

[0012] In some embodiments, the method further includes: determining a feature-level loss function based on the multiple target discrete vectors and the second image features; optimizing the second encoder based on the feature-level loss function to obtain a target encoder; determining the trained target image restoration model based on the target discrete codebook and the target prediction module includes: determining the target image restoration model based on the target discrete codebook, the target prediction module and the target encoder.

[0013] In some embodiments, the above method also includes: decoding the above multiple matching discrete vectors based on the decoder to obtain the first repaired image corresponding to the above first image; optimizing the above decoder based on the above first image and the above first repaired image to obtain the target decoder; the above determination of the trained target image repair model based on the above target discrete codebook and the above target prediction module includes: determining the above target image repair model based on the above target discrete codebook, the above target prediction module and the above target decoder.

[0014] In some embodiments, the method further includes: decoding the plurality of predicted discrete vectors based on the target decoder to obtain a second restored image corresponding to the second image.

[0015] In some embodiments, the above-mentioned decoding processing of the above-mentioned multiple predicted discrete vectors based on the above-mentioned target decoder to obtain the second restored image corresponding to the above-mentioned second image includes: determining multiple enhanced predicted discrete vectors based on the above-mentioned second image features and the above-mentioned multiple predicted discrete vectors; decoding processing of the above-mentioned multiple enhanced predicted discrete vectors based on the above-mentioned target decoder to obtain the above-mentioned second restored image.

[0016] A second aspect of an embodiment of the present application provides an image restoration method, comprising: encoding a target image based on a target encoder in a target image restoration model to obtain target image features corresponding to the target image; determining multiple target prediction discrete vectors corresponding to the target image features in a target discrete codebook based on a target prediction module in the target image restoration model; wherein the target discrete codebook includes multiple target discrete vectors, and the target discrete vectors include the target prediction discrete vector; decoding the multiple target prediction discrete vectors based on a target decoder in the target image restoration model to obtain a target restoration image corresponding to the target image.

[0017] In some embodiments, the target prediction module in the target image restoration model determines multiple target prediction discrete vectors corresponding to the target image features in the target discrete codebook, including: determining multiple target prediction vector indexes corresponding to the target image features in the target discrete codebook based on the target prediction module; wherein the target discrete codebook includes the target vector index corresponding to each of the multiple target discrete vectors, and the target vector index includes the target prediction vector index; determining the multiple target prediction discrete vectors corresponding to the target image features in the target discrete codebook based on the multiple target prediction vector indexes.

[0018] In some embodiments, the target decoder in the target image restoration model decodes the multiple target prediction discrete vectors to obtain a target restoration image corresponding to the target image, including: determining multiple enhanced target prediction discrete vectors based on the target image features and the multiple target prediction discrete vectors; and decoding the multiple enhanced target prediction discrete vectors based on the target decoder to obtain the target restoration image.

[0019] A third aspect of an embodiment of the present application provides a training device for an image restoration model, comprising: a first optimization module, configured to determine a first image feature corresponding to a first image based on the image restoration model, and optimize multiple preset discrete vectors in a discrete codebook based on the first image feature to obtain a target discrete codebook; wherein the target discrete codebook includes multiple target discrete vectors; a determination module, configured to determine a second image feature corresponding to a second image based on the image restoration model, and determine multiple predicted discrete vectors corresponding to the second image feature in the target discrete codebook based on a prediction module in the image restoration model; wherein a quality parameter of the first image is higher than a quality parameter of the second image, and the target discrete vector includes the predicted discrete vector; a second optimization module, configured to optimize the prediction module based on the multiple target discrete vectors and the multiple predicted discrete vectors to obtain a target prediction module, and determine a trained target image restoration model based on the target discrete codebook and the target prediction module.

[0020] A fourth aspect of an embodiment of the present application provides an image restoration device, comprising: an encoding module, configured to encode a target image based on a target encoder in a target image restoration model to obtain target image features corresponding to the target image; a prediction module, configured to determine, in a target discrete codebook, a plurality of target prediction discrete vectors corresponding to the target image features based on a target prediction module in the target image restoration model; wherein the target discrete codebook includes a plurality of target discrete vectors, and the target discrete vectors include the target prediction discrete vector; a decoding module, configured to decode the plurality of target prediction discrete vectors based on a target decoder in the target image restoration model to obtain a target restoration image corresponding to the target image.

[0021] The fifth aspect of an embodiment of the present application provides an electronic device, comprising: a processor and a memory, the memory being used to store computer executable instructions; the processor being used to read the instructions from the memory and execute the instructions to implement the image restoration model training method described in the first aspect or the image restoration method described in the second aspect.

[0022] The sixth aspect of an embodiment of the present application provides a computer-readable storage medium, which stores computer program instructions. When the computer reads the instructions, it executes the image restoration model training method described in the first aspect or the image restoration method described in the second aspect.

[0023] A seventh aspect of an embodiment of the present application provides a computer program product, which includes a computer program stored on a non-transitory computer-readable storage medium, and the computer program includes program instructions. When the program instructions are executed by a computer, the computer executes the training method of the image restoration model described in the first aspect or the image restoration method described in the second aspect.

[0024] The training method of the image restoration model provided in the embodiment of the present application can train the discrete codebook in the image restoration model based on the first image, and the trained target discrete codebook includes high-quality discrete vectors (i.e., target discrete vectors); and train the prediction module in the image restoration model based on the second image (wherein the quality parameter of the first image is higher than the quality parameter of the second image) and the target discrete vector in the target discrete codebook to obtain a trained target prediction module, and determine the trained target image restoration model based on the target discrete codebook and the target prediction module. The present application trains the target discrete codebook to store high-quality features in the image, and restores low-quality images based on the target discrete codebook through the target prediction module, thereby eliminating the need to rely on prior information, thereby improving the quality of image restoration, and enhancing the robustness and generalization of image restoration in image degradation scenarios. BRIEF DESCRIPTION OF THE DRAWINGS

[0025] In order to more clearly illustrate the technical solutions of the embodiments of the present invention, the drawings required for use in the embodiments will be briefly introduced below. Obviously, the drawings described below are only some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on these drawings without paying creative work.

[0026] Figure 1 A schematic diagram of an image restoration model provided in some embodiments of the present application;

[0027] Figure 2 A flowchart of a method for training an image restoration model provided in some embodiments of the present application;

[0028] Figure 3 A schematic diagram of a training method for an image restoration model provided in some embodiments of the present application;

[0029] Figure 4 A flowchart of another method for training an image restoration model provided in some embodiments of the present application;

[0030] Figure 5 A flowchart of a training method for another image restoration model provided in some embodiments of the present application;

[0031] Figure 6A flowchart of another method for training an image restoration model provided in some embodiments of the present application;

[0032] Figure 7 A flowchart of an image restoration method provided in some embodiments of the present application;

[0033] Figure 8 A schematic diagram of an image restoration method provided in some embodiments of the present application;

[0034] Fig. 9 A schematic diagram of a training device for an image restoration model provided in some embodiments of the present application;

[0035] Fig.10 A schematic diagram of an image restoration device provided in some embodiments of the present application;

[0036] Fig.11 A schematic diagram of an electronic device according to some embodiments of the present application. DETAILED DESCRIPTION

[0037] In order to enable those skilled in the art to better understand the technical solutions in the embodiments of the present invention and to make the above-mentioned purposes, features and advantages of the embodiments of the present invention more obvious and understandable, the technical solutions in the embodiments of the present invention are further described in detail below with reference to the accompanying drawings.

[0038] The training method of the image restoration model provided in this application is described in detail below in conjunction with the accompanying drawings.

[0039] Since there is no way to accurately know the cause of image degradation, a low-quality input image (i.e., a degraded image) may correspond to countless potential high-quality output images (i.e., the restored image). Therefore, existing image restoration methods usually rely on prior information to solve this problem, where the prior information includes geometric prior, reference prior, and generation prior. However, image restoration methods based on prior information usually have some shortcomings. Taking face images as an example, image restoration methods based on geometric priors (such as PSFRGAN) use facial analysis maps as prior information to assist face restoration, but this method lacks the ability to restore image details; image restoration methods based on reference priors (such as DFDNet) use facial component dictionaries as prior information, but only contain specific facial components (such as eyes, mouth and nose), resulting in a lack of fine details on the face in the repaired image, and have high computational complexity and long training time, which will affect the efficiency of actual applications; image restoration methods based on generative priors (such as GPEN and GFP-GAN) use pre-trained generative models (such as StyleGAN2) as prior information, but when dealing with real scenes with severe degradation and diverse degradation, they may perform poorly due to the limitations of pre-trained models.

[0040] In summary, the above-mentioned image restoration method based on prior information projects the degraded face into a continuous latent space through iterative latent optimization or direct latent coding. Although it can improve the face restoration effect to a certain extent, this image restoration method still has some problems, such as: high sensitivity to degradation, that is, when the quality of the input low-quality image decreases, the quality of the output restoration image will decrease significantly; the expressive power of prior information is limited; excessive reliance on low-quality input image information leads to uncertainty in mapping from low-quality images to high-quality images; it is difficult to maintain high fidelity when repairing severely degraded images; the restored images often lack fine-grained facial details and may contain obvious artifacts.

[0041] Based on the above technical problems, the present application provides an image restoration model training method, an image restoration method and a device, which can improve the effect of image restoration and enhance the robustness and generalization ability of image restoration in image degradation scenarios.

[0042] Figure 1 A schematic diagram of an image restoration model provided in some embodiments of the present application. Figure 1 As shown, the image restoration model 100 to be trained includes a first encoder 110, a discrete codebook 120, a decoder 130, a second encoder 140 and a prediction module 150. The image restoration model 100 provided in the embodiment of the present application can process images containing any content, such as images containing faces, objects, landscapes, etc. In the following embodiments, face images are used as an example for exemplary description.

[0043] Among them, the first encoder 110 and the second encoder 140 are used to encode the images (such as the first image, the second image, etc.) input into the image restoration model 100 to obtain corresponding image features (such as the first image features corresponding to the first image, the second image features corresponding to the second image, etc.).

[0044] The discrete codebook 120 is used to store a plurality of preset discrete vectors, each of which corresponds to a preset vector index. The preset discrete vectors are obtained by hierarchically quantizing the high-dimensional features corresponding to the image, that is, each preset discrete vector can correspond to features at different levels in the image, such as the composition, texture, color distribution, semantic information, etc. of the image. For example, the discrete codebook 120 can store quantitative features such as facial components, facial details, and semantic information in a face image, such as features corresponding to components such as hair, eyes, nose, and mouth, and features corresponding to facial details such as texture and color distribution, as well as semantic information such as the size and shape of each facial component, facial expressions, facial accessories, facial postures, gender, and age of the person.

[0045] The decoder 130 is used to decode multiple discrete vectors (such as multiple matching discrete vectors, multiple predicted discrete vectors, etc.) and generate corresponding images (such as a first restored image and a second restored image).

[0046] The prediction module 150 is used to predict the discrete vector corresponding to the second image feature of the second image based on multiple target discrete vectors in the trained target discrete codebook to obtain a predicted discrete vector, so that the decoder 130 generates a second restored image corresponding to the second image based on the predicted discrete vector.

[0047] The following is an exemplary description of the training method of the image restoration model 100.

[0048] Figure 2 A flowchart of a method for training an image restoration model provided in some embodiments of the present application. Figure 2 As shown, the training method of the image restoration model may include steps 210 to 230.

[0049] Step 210: determine a first image feature corresponding to the first image based on the image restoration model, and optimize a plurality of preset discrete vectors in a discrete codebook based on the first image feature to obtain a target discrete codebook.

[0050] refer to Figure 3 A schematic diagram of a training method for an image restoration model is shown in FIG. Figure 3 As shown, the training process of the image restoration model can be divided into two stages. The main purpose of the first stage is to train the discrete codebook, and the main purpose of the second stage is to train the prediction module (such as the transformer). The following first illustrates the training process of the first stage.

[0051] In some embodiments, in the first stage, the image restoration model obtains an input first image I h (I h ∈R H ×W×3 ), the first image may be an image with a higher quality parameter (i.e., a high-quality image), that is, an image with a lower degree of degradation from which high-quality features can be extracted.

[0052] Exemplarily, in the above step 210, "determining the first image feature corresponding to the first image based on the image restoration model" may include: encoding the first image based on the first encoder in the image restoration model to obtain the first image feature and multiple first image sub-features corresponding to the first image feature.

[0053] In some embodiments, after the first image is acquired, the first encoder E in the image restoration model can be used to HThe first image is encoded to map the first image to a latent space and obtain a first image feature Z corresponding to the first image. h (Z h ∈R m×n×d ), namely Z h =E H (I h ). Further, the first image feature Z can be transformed into h Mapped to the dimension of (m·n)×d, the first image feature Z h Flattened into m·n feature vectors, thus obtaining the first image feature Z h Corresponding multiple first image sub-features Z i , that is, the first image feature Z h =[Z1,Z2,Z3…Z i ]. The value range of i is 1 to m·n.

[0054] Figure 4 A flowchart of another method for training an image restoration model provided in some embodiments of the present application. Figure 4 As shown, in the above step 210 , “optimizing multiple preset discrete vectors in the discrete codebook based on the first image feature to obtain a target discrete codebook” includes steps 410 to 420 .

[0055] Step 410: Determine a plurality of matching discrete vectors corresponding to the first image feature in a discrete codebook.

[0056] In some embodiments, after determining the first image feature corresponding to the first image and the multiple first image sub-features corresponding to the first image feature, the first image sub-features may be quantized (i.e., discretized) based on a discrete codebook to determine multiple matching discrete vectors corresponding to the first image sub-features. The preset discrete vectors in the discrete codebook include matching discrete vectors.

[0057] Exemplarily, the above-mentioned determining multiple matching discrete vectors corresponding to the first image feature in the discrete codebook includes: determining a candidate similarity between the first image sub-feature and each preset discrete vector in a plurality of preset discrete vectors; determining a target similarity with the highest similarity among the multiple candidate similarities, and determining a target preset discrete vector corresponding to the target similarity as a matching discrete vector corresponding to the first image sub-feature; wherein the preset discrete vector includes a target preset discrete vector.

[0058] In some embodiments, the discrete codebook includes a plurality of preset discrete vectors and a preset vector index corresponding to each preset discrete vector. k )} k∈[K] , where ck is the preset discrete vector, k is the preset vector index corresponding to the preset discrete vector, and the discrete codebook C includes a total of K preset discrete vectors c k .

[0059] Exemplarily, the size of the discrete codebook (i.e., the number of preset discrete vectors included) can be set according to actual conditions. The size of the discrete codebook is related to the degree of compression of the information (i.e., the compression ratio). The larger the size of the discrete codebook, the lower the accuracy of quantizing image features and the smaller the compression ratio; the smaller the size of the discrete codebook, the higher the accuracy of quantizing image features and the higher the compression ratio. Therefore, when setting the size of the discrete codebook, a more appropriate size can be selected so that the discrete vectors in the discrete codebook are more compact, so that the subsequent prediction module is more efficient when predicting based on the discrete codebook, and the prediction effect of the prediction module can also be ensured.

[0060] After determining the plurality of first image sub-features corresponding to the first image feature, the mapping functions Q of different feature layers can be used. l(·,·) The first image sub-feature Z is calculated based on the preset discrete vector in the discrete codebook. i Quantization is performed to determine the target preset discrete vector with the highest similarity to each first image sub-feature among the multiple preset discrete vectors.

[0061] like Figure 3 As shown, the first image sub-feature can be matched with each preset discrete vector in a plurality of preset discrete vectors to determine the candidate similarity between the first image sub-feature and each preset discrete vector. The similarity matching can be implemented by nearest neighbor matching, and the nearest neighbor matching can be implemented based on distance formulas such as Euclidean distance and Manhattan distance.

[0062] After determining the candidate similarities, the target preset discrete vector corresponding to the target similarity with the highest similarity among the candidate similarities can be determined as the matching discrete vector corresponding to the first image sub-feature, and the matching vector index corresponding to the matching discrete vector can be determined. That is, the matching discrete vector is the discrete vector with the highest similarity to the first image sub-feature in the discrete codebook.

[0063] Among them, multiple matching vector indexes s i A corresponding matching index sequence S can be formed, that is, S = [s1, s2, s3…s i ], the value range of i is 1 to m·n.

[0064] For example, the discrete codebook C includes three preset discrete vectors c1 = [1.0, 0.0, 0.0], c2 = [0.0, 1.0, 0.0], c3 = [0.0, 0.0, 1.0], and the first image feature Z h Taking four first image sub-features Z1 = [0.1, 0.9, 0.1], Z2 = [0.3, 0.1, 0.7], Z3 = [0.8, 0.1, 0.1], and Z4 = [0.2, 0.8, 0.2] as an example, the candidate similarities between each first image sub-feature and each preset discrete vector calculated based on the Euclidean distance are as follows:

[0065] distance(Z1,c1)=1.62, distance(Z1,c2)=0.02, distance(Z1,c3)=1.82;

[0066] distance(Z2,c1)=1.26, distance(Z2,c2)=1.26, distance(Z2,c3)=0.26;

[0067] distance(Z3,c1)=0.02, distance(Z3,c2)=0.02, distance(Z3,c3)=1.62;

[0068] distance(Z4,c1)=1.32, distance(Z4,c2)=0.32, distance(Z4,c3)=1.32;

[0069] It can be seen that for the first image sub-feature Z1, the target preset discrete vector with the highest similarity (i.e., the smallest Euclidean distance) among all preset discrete vectors is c2; for the first image sub-feature Z2, the target preset discrete vector with the highest similarity among all preset discrete vectors is c3; for the first image sub-feature Z3, the target preset discrete vector with the highest similarity among all preset discrete vectors is c1 (or c2); for the first image sub-feature Z4, the target preset discrete vector with the highest similarity among all preset discrete vectors is c2. That is, the first image feature Z h The corresponding multiple matching discrete vectors are c2, c3, c1, c2, so the first image feature Z h The corresponding first quantitative feature Z c =[c2, c3, c1, c2], wherein the first quantitative feature is composed of a plurality of matching discrete vectors; the first quantitative feature Z c The corresponding matching index sequence S = [2, 3, 1, 2].

[0070] Exemplarily, after determining multiple matching discrete vectors corresponding to the first image features, the multiple matching discrete vectors can be decoded based on the decoder to obtain a first restored image corresponding to the first image; and the decoder can be optimized based on the first image and the first restored image to obtain a target decoder.

[0071] like Figure 3 As shown, after mapping the first image feature into a plurality of matching discrete vectors (i.e., the first quantized feature) based on the discrete codebook, the decoder D H For the first quantitative feature Z c Decoding is performed to convert the first quantized feature Z c Mapped to the image input space and generating a first repaired image corresponding to the first image, i.e., the first repaired image I rec =D H (Z c ).

[0072] Step 420: Optimize multiple preset discrete vectors in a discrete codebook based on the first image feature and multiple matching discrete vectors to obtain a target discrete codebook.

[0073] In some embodiments, after determining the above-mentioned multiple matching discrete vectors (i.e., the first quantitative features), multiple preset discrete vectors in the discrete codebook can be optimized based on the first image features corresponding to the first image and the multiple matching discrete vectors, so as to obtain a target discrete codebook C′. The target discrete codebook includes multiple target discrete vectors (e.g., K target discrete vectors).

[0074] Exemplarily, a quantization loss function can be determined based on the first image feature and the first quantization feature, and multiple preset discrete vectors in the discrete codebook can be optimized based on the quantization loss function. Since the quantized differential is indefinite, a straight through estimator (STE) can be used for gradient approximation in the reverse calculation process.

[0075] In some examples, the above quantization loss function can be determined according to formula (1):

[0076]

[0077] Among them, L Q is the quantitative loss function, Z h is the first image feature, Z c is the first quantitative feature, β is a hyperparameter (β∈[0,1]), and sg[·] is a stop gradient operator used to prevent gradient backpropagation. Q The smaller the value, the higher the Z c With Z h The higher the similarity between them.

[0078] Through the above scheme, based on the encoding result of the first image by the first encoder (i.e., the first image feature) and the first quantized feature obtained by similarity matching in the discrete codebook, multiple preset discrete vectors in the discrete codebook are updated and optimized by gradient descent, so that the multiple preset discrete vectors in the discrete codebook can be closer to the encoding result of the first encoder (i.e., the first image feature). The final trained target discrete codebook includes multiple high-quality discrete vectors (i.e., multiple target discrete vectors).

[0079] In some embodiments, after determining the first repaired image, the decoder may be optimized based on the first image and the first repaired image, thereby obtaining a target decoder D H ′.

[0080] Exemplarily, a mean square error loss function may be determined based on the first image and the first restored image, and the decoder may be optimized based on the mean square error loss function. In some examples, the mean square error loss function may be determined according to formula (2):

[0081]

[0082] Among them, L MSE is the mean square error loss function, I h is the first image, I rec is the first restored image, M is the first image I h The total number of pixels of the first repaired image I rec The total number of pixels. MSE The smaller the value, the better the first image I h With the first repaired image I rec The higher the similarity between them.

[0083] Through the above scheme, the decoder is optimized and trained based on the first image and the first repaired image, thereby improving the decoder's ability to repair images, so that the trained target decoder can generate a higher quality repaired image based on the quantitative features corresponding to the image to be repaired.

[0084] In addition, the first encoder can also be optimized and trained based on the above-mentioned mean square error loss function, so that the trained first encoder can generate more effective image features when encoding the image; further, the discrete codebook can be optimized based on the more effective image features, so that the target discrete vectors stored in the target discrete codebook can be more accurate and effective, and the image restoration capability of the target decoder can be stronger.

[0085] Furthermore, the first encoder, the discrete codebook, and the decoder can be optimized and trained simultaneously based on the first loss function, wherein the first loss function can be determined according to the sum of the quantization loss function and the mean square error loss function. In some examples, the first loss function L can be determined according to formula (3): HVQ :

[0086] L HVQ =L MSE (I h ,I rec )+L Q ; (3)

[0087] Through the above scheme, in the first stage of training of the image reconstruction model, the target discrete codebook and the target decoder are obtained through self-reconstruction learning training to store the high-quality visual part of the image based on the target discrete codebook, and generate a high-quality repaired image based on the target decoder. The input high-quality image features are decomposed into multi-level discrete features, and these features are mapped to the discrete codebook, so that the high-quality image features can be represented by the discrete vectors in the discrete codebook. Therefore, the target discrete vectors in the trained target discrete codebook not only contain the global structural features of the image, but also contain fine-grained image features, so as to provide rich feature information for the subsequent image restoration process to improve the effect of image restoration.

[0088] After obtaining the target discrete codebook and target decoder through the above-mentioned first-stage training process, the image reconstruction model can continue to be optimized in the second stage. It can be understood that when the image with low input quality parameters (i.e., low-quality images) is repaired by the image restoration model, the image features extracted from the low-quality images are usually incomplete or inaccurate due to the high degree of degradation of the low-quality images. Therefore, if the discrete vector corresponding to the image feature is matched in the target discrete codebook directly by the nearest neighbor matching method, an erroneous matching result may be generated; and when the distance between the features of the low-quality image and each discrete vector in the discrete codebook is large, it will be difficult to match the accurate discrete vector, thereby affecting the effect of image restoration.

[0089] Therefore, in the embodiment of the present application, the discrete vectors corresponding to the features of the low-quality image are predicted by training the prediction module in the second stage, so as to improve the accuracy of the image reconstruction model in determining multiple discrete vectors matching the features of the low-quality image, thereby improving the effect of image restoration. The following is an exemplary description of the training process of the second stage.

[0090] Step 220: determine a second image feature corresponding to the second image based on the image restoration model, and determine a plurality of predicted discrete vectors corresponding to the second image feature in a target discrete codebook based on a prediction module in the image restoration model.

[0091] like Figure 3 As shown, in the second stage, the image restoration model first obtains a second image I with lower quality parameters l (i.e., a low-quality image). Among them, the quality parameter of the first image is higher than the quality parameter of the second image. That is, the second image is an image with a high degree of degradation and it is difficult to extract high-quality features from it. Then, the second encoder E in the image restoration model can be used to L The second image is encoded to obtain a low-quality second image feature Z corresponding to the second image. l (Z l ∈R m×n×d ), namely Z l =E L (I l ).

[0092] For example, the second encoder E L The first encoder E H After fine-tuning, for example, the first encoder E H The weight parameters in are adjusted so that the resulting second encoder E L Can adapt to low-quality image input.

[0093] Furthermore, the second image feature Z can be transformed into l Mapped to the dimension of (m·n)×d, the second image feature Z l Flattened into m·n feature vectors, thus obtaining the second image feature Z l Corresponding multiple second image sub-features Z v , that is, the second image feature Z l =[Z1,Z2,Z3…Z v ]. The value range of v is 1 to m·n.

[0094] After obtaining the second image feature and the corresponding multiple second image sub-features, multiple predicted discrete vectors corresponding to the second image feature can be determined based on the target discrete codebook obtained by the prediction module (such as Transformer) after the first stage of training. The target discrete codebook includes multiple target discrete vectors, and the target discrete vector includes the predicted discrete vector.

[0095] Figure 5 A flowchart of another method for training an image restoration model provided in some embodiments of the present application. Figure 5As shown, in the above step 220, "determining multiple predicted discrete vectors corresponding to the second image feature in the target discrete codebook based on the prediction module in the image restoration model" includes steps 510 to 520.

[0096] Step 510: Determine multiple prediction vector indexes corresponding to the second image feature in the target discrete codebook based on the prediction module.

[0097] For example, when the prediction module determines multiple prediction discrete vectors corresponding to the second image feature in the target discrete codebook, multiple prediction vector indexes corresponding to the second image feature can be first determined in the target discrete codebook, and then the corresponding multiple prediction discrete vectors can be determined according to the multiple prediction vector indexes, thereby reducing the complexity of the prediction module when performing prediction and significantly improving the prediction efficiency. k ′ corresponds to the target vector index k′, which includes the predicted vector index

[0098] Exemplarily, the above step 510 may include: determining a high-level feature corresponding to the second image feature based on a prediction module; wherein the high-level feature is determined using a self-attention mechanism; and determining multiple prediction vector indexes corresponding to the high-level feature in a target discrete codebook based on the prediction module.

[0099] In some embodiments, when the prediction module determines multiple prediction vector indexes corresponding to the second image feature in the target discrete codebook, it can first use a self-attention mechanism to analyze the multiple second image sub-features Z. v Perform global modeling, feature enhancement, and semantic extraction to obtain high-level features corresponding to the second image features. That is, the high-level features are the original low-quality features (i.e., multiple second image sub-features Z v ) is a result of multiple enhancements and fusions, which contains local information, global context information, high-level semantic information, etc. corresponding to the second image, so that the image restoration model can obtain more accurate prediction results when predicting discrete vectors based on the high-level features, thereby improving the effect of image restoration.

[0100] In some examples, the calculation formula of the Sth self-attention block of the prediction module can refer to formula (4):

[0101] X s+1 =Softmax(Q s K s )V s +X s ; (4)

[0102] Among them, the input of the first self-attention block is X0, which is the second image sub-feature Z v ; The input feature of the Sth self-attention block is X s ;Q s , K s 、V s They are respectively passed through the linear layer according to the input feature X s The query, key, and value obtained; the output feature of the Sth self-attention block is X s+1 , X s+1 That is, the high-level feature corresponding to the second image feature.

[0103] After obtaining the high-level feature X corresponding to the second image feature s+1 Afterwards, the high-level features X can be transformed into s+1 Project to a preset dimension (such as (m·n)×N), so as to output the multi-dimensional high-level feature X s+1 Compress or expand to a dimension suitable for image restoration model processing. Among them, N can be determined according to the size of the discrete codebook. For example, if there are K target discrete vectors in the target discrete codebook, the high-level feature C s+1 Projected to (m·n)×K dimensions.

[0104] Next, the prediction module can predict multiple vector indexes corresponding to the projected high-level features in the target discrete codebook C′ to obtain multiple predicted vector indexes And get multiple prediction vector indexes The prediction vector index sequence composed of Right now

[0105] Step 520: Determine a plurality of prediction discrete vectors corresponding to the second image feature in the target discrete codebook based on the plurality of prediction vector indices.

[0106] like Figure 3 As shown, after obtaining multiple prediction vector indexes After that, the image restoration model can further determine the index of each prediction vector in the target discrete codebook C′ The corresponding prediction discrete vector Thus, we get a discrete vector consisting of multiple predictions The second quantitative feature C c (C c ∈R m×n×d ). Among them, the target discrete vector c k ′ includes the predicted discrete vector

[0107] In some embodiments, after determining the plurality of predicted discrete vectors (i.e., second quantized features) corresponding to the second image features, the second quantized features may be decoded based on the target decoder to obtain a second restored image corresponding to the second image, wherein the quality parameter of the second restored image is higher than that of the second image, i.e., the degradation degree of the second restored image is lower than that of the second image.

[0108] In some examples, when determining the second inpainted image corresponding to the second image, in order to further enhance the fine-grained features in the second quantized features, a gated cross-attention mechanism (Gated Context Attention) can be used to perform feature enhancement processing on the second quantized features. Exemplarily, the gated cross-attention mechanism can be used to process the second image features and multiple predicted discrete vectors (i.e., the second quantized features) to determine multiple enhanced predicted discrete vectors (i.e., enhanced second quantized features); then, based on the target decoder, the multiple enhanced predicted discrete vectors (i.e., enhanced second quantized features) are decoded to obtain the second inpainted image.

[0109] In some examples, the method of using the gated cross attention mechanism to determine the enhanced second quantitative feature can be shown in formula (5):

[0110]

[0111] Among them, Z L is the second image feature, C c is the second quantitative feature, It is the enhanced second quantitative feature.

[0112] Step 230, optimizing the prediction module based on multiple target discrete vectors and multiple prediction discrete vectors to obtain a target prediction module, and determining a trained target image restoration model based on the target discrete codebook and the target prediction module.

[0113] After determining the above-mentioned multiple target discrete vectors and multiple predicted discrete vectors, a second loss function can be determined based on the multiple target discrete vectors and the multiple predicted discrete vectors, and the prediction module can be optimized based on the second loss function to obtain a target prediction module.

[0114] Figure 6 A flowchart of another method for training an image restoration model provided in some embodiments of the present application. Figure 6 As shown, the image restoration method may include steps 610 to 620.

[0115] Step 610: Determine a code-level loss function based on a plurality of target vector indices corresponding to a plurality of target discrete vectors and a plurality of prediction vector indices corresponding to a plurality of prediction discrete vectors.

[0116] In some embodiments, the code level loss function may be determined based on multiple target vector indices and multiple prediction vector indices in the target discrete codebook. In some examples, the code level loss function may be determined according to formula (6):

[0117]

[0118] in, is the code-level loss function, s i is the target vector index, is the prediction vector index.

[0119] Step 620, optimizing the prediction module based on the code level loss function to obtain a target prediction module.

[0120] By determining a code-level loss function based on multiple target vector indices and multiple prediction vector indices in a target discrete codebook and optimizing the prediction module based on the code-level loss function, the accuracy of the prediction module in predicting discrete vectors corresponding to low-quality image features can be improved, thereby improving the effect of the image restoration model in restoring the image.

[0121] In some embodiments, the second encoder may be optimized based on multiple target discrete vectors and second image features in the target discrete codebook to obtain a trained second encoder (hereinafter referred to as a target encoder). Exemplarily, a feature-level loss function may be determined based on multiple target discrete vectors and second image features; and the second encoder may be optimized based on the feature-level loss function to obtain a target encoder.

[0122] In some examples, the feature-level loss function can be determined according to formula (7):

[0123]

[0124] in, is the feature-level loss function, Z l is the second image feature, c k ′ is the target discrete vector in the target discrete codebook.

[0125] Through the above scheme, the second encoder is optimized based on multiple target discrete vectors and second image features in the target discrete codebook, which can improve the quality of image features generated by the second encoder when encoding low-quality images, thereby improving the effect of the image restoration model when restoring the image.

[0126] Furthermore, the prediction module and the second encoder may be trained simultaneously based on a second loss function, wherein the second loss function may be determined based on the sum of the code-level loss function and the feature-level loss function. In some examples, the second loss function L may be determined according to formula (8): tf :

[0127]

[0128] Among them, λ token is a hyperparameter.

[0129] Through the above scheme, the second stage of optimization training of the image restoration model can improve the accuracy of the prediction module in predicting the discrete vector corresponding to the low-quality image features, and improve the quality of the image features generated when the second encoder encodes the low-quality image, thereby further improving the effect of the image restoration model on the image.

[0130] After completing the training of the image restoration model through the first and second stages and obtaining the target image restoration model (including the target encoder, the target discrete codebook, the target prediction module and the target decoder), the input target image can be restored based on the target image restoration model to obtain and output the target restoration image corresponding to the target image. The following is an exemplary description of the image restoration method based on the target image restoration model.

[0131] Figure 7 A flowchart of an image restoration method provided in some embodiments of the present application. Figure 7 As shown, the image restoration method may include steps 710 to 730.

[0132] Step 710, encoding the target image based on the target encoder in the target image restoration model to obtain target image features corresponding to the target image.

[0133] refer to Figure 8 A schematic diagram of an image restoration method is shown in FIG. Figure 8 As shown, the target image restoration model obtains the input target image. In the embodiment of the present application, the quality parameter (i.e., the degree of degradation) of the target image is not limited. Then, the target encoder E in the target image restoration model is used. L 'Encode the target image to obtain the target image features corresponding to the target image.

[0134] It can be understood that the implementation of step 710 can refer to the description of step 220, which will not be repeated here.

[0135] Step 720: Based on the target prediction module in the target image restoration model, a plurality of target prediction discrete vectors corresponding to the target image features are determined in the target discrete codebook.

[0136] Exemplarily, after obtaining the target image features, the target prediction module in the target image restoration model can be used to determine multiple target prediction discrete vectors corresponding to the target image features in the target discrete codebook C′; multiple target prediction vector indexes corresponding to the target image features can also be determined in the target discrete codebook C′, and the target prediction discrete vectors corresponding to each target prediction vector index can be determined in the target discrete codebook C′. The target discrete codebook includes multiple target discrete vectors, the target discrete vector includes a target prediction discrete vector; the target discrete codebook includes a target vector index corresponding to each target discrete vector in the multiple target discrete vectors, and the target vector index includes a target prediction vector index.

[0137] Exemplarily, the target quantitative feature corresponding to the target image feature may be determined based on a plurality of target prediction discrete vectors corresponding to the target image feature, wherein the target quantitative feature is composed of a plurality of target prediction discrete vectors.

[0138] It can be understood that the implementation of step 720 can refer to the description of step 220 and steps 510 to 520, and will not be repeated here.

[0139] Step 730 , based on the target decoder in the target image restoration model, a plurality of target prediction discrete vectors are decoded to obtain a target restoration image corresponding to the target image.

[0140] For example, after obtaining multiple target prediction discrete vectors (i.e., target quantized features) corresponding to the target image features, a gated cross-attention mechanism can be used to determine multiple enhanced target prediction discrete vectors (i.e., enhanced target quantized features) based on the target image features and the multiple target prediction discrete vectors; then, based on the target decoder D H 'Decoding multiple enhanced target prediction discrete vectors (i.e., enhanced target quantized features) to obtain a target restoration image corresponding to the target image. The quality parameter of the target restoration image is higher than that of the input target image, i.e., the degradation degree of the target restoration image is lower than that of the target image.

[0141] It can be understood that the implementation method of step 730 can refer to the description of the above-mentioned embodiment of "after determining multiple predicted discrete vectors (i.e., second quantized features) corresponding to the second image features, the second quantized features can be decoded based on the target decoder to obtain a second restored image corresponding to the second image", which will not be repeated here.

[0142] Through the above scheme, in the process of image restoration processing of the input target image based on the trained target image restoration model, the target image is restored through the global image features and local image features stored in the target discrete codebook, thereby improving the ability to restore the overall features and detail features of the image; and through the prediction module, the global image structure is effectively modeled, and multiple target prediction discrete vectors corresponding to the target image are predicted, thereby enhancing the robustness of the image restoration model in image restoration under degraded scenarios; and the gated cross-attention mechanism is used to compensate for the information that may be lost in the quantization process, thereby further improving the quality of image restoration.

[0143] Fig. 9 A schematic diagram of a training device for an image restoration model provided in some embodiments of the present application. Fig. 9 As shown, the training device 900 for the image restoration model includes a first optimization module 910 , a determination module 920 and a second optimization module 930 .

[0144] The first optimization module 910 is configured to determine a first image feature corresponding to the first image based on the image restoration model, and optimize a plurality of preset discrete vectors in the discrete codebook based on the first image feature to obtain a target discrete codebook.

[0145] The target discrete codebook includes multiple target discrete vectors.

[0146] The determination module 920 is configured to determine the second image feature corresponding to the second image based on the image restoration model, and determine multiple predicted discrete vectors corresponding to the second image feature in the target discrete codebook based on the prediction module in the image restoration model.

[0147] The quality parameter of the first image is higher than the quality parameter of the second image, and the target discrete vector includes a predicted discrete vector.

[0148] The second optimization module 930 is configured to optimize the prediction module based on multiple target discrete vectors and multiple prediction discrete vectors to obtain a target prediction module, and determine a trained target image restoration model based on the target discrete codebook and the target prediction module.

[0149] In some embodiments, the first optimization module 910 is specifically configured to: determine multiple matching discrete vectors corresponding to the first image feature in the discrete codebook; wherein the preset discrete vector includes the matching discrete vector; based on the first image feature and the multiple matching discrete vectors, optimize the multiple preset discrete vectors in the discrete codebook to obtain a target discrete codebook.

[0150] In some embodiments, the first optimization module 910 is specifically configured to: encode the first image based on the first encoder in the image restoration model to obtain a first image feature and multiple first image sub-features corresponding to the first image feature; determine the candidate similarities between the first image sub-feature and each preset discrete vector in a plurality of preset discrete vectors; determine the target similarity with the highest similarity among the multiple candidate similarities, and determine the target preset discrete vector corresponding to the target similarity as the matching discrete vector corresponding to the first image sub-feature; wherein the preset discrete vector includes the target preset discrete vector.

[0151] In some embodiments, the determination module 920 is specifically configured to: determine multiple prediction vector indices corresponding to the second image feature in the target discrete codebook based on the prediction module; wherein the target discrete codebook includes target vector indices corresponding to each target discrete vector, and the target vector index includes a prediction vector index; determine multiple prediction discrete vectors corresponding to the second image feature in the target discrete codebook based on multiple prediction vector indices.

[0152] In some embodiments, the determination module 920 is specifically configured to: determine the high-level features corresponding to the second image features based on the prediction module; wherein the high-level features are determined using a self-attention mechanism; and determine multiple prediction vector indexes corresponding to the high-level features in the target discrete codebook based on the prediction module.

[0153] In some embodiments, the second optimization module 930 is specifically configured to: determine a code-level loss function based on multiple target vector indices corresponding to multiple target discrete vectors and multiple prediction vector indices corresponding to multiple prediction discrete vectors; optimize the prediction module based on the code-level loss function to obtain a target prediction module.

[0154] In some embodiments, the determination module 920 is specifically configured to: perform encoding processing on the second image based on the second encoder in the image restoration model to obtain a second image feature corresponding to the second image.

[0155] like Fig. 9 As shown, the image restoration model training device 900 also includes a third optimization module 940.

[0156] In some embodiments, the third optimization module 940 is configured to: determine a feature-level loss function based on multiple target discrete vectors and second image features; optimize the second encoder based on the feature-level loss function to obtain a target encoder; the second optimization module 930 is specifically configured to: determine a target image restoration model based on a target discrete codebook, a target prediction module and a target encoder.

[0157] like Fig. 9 As shown, the image restoration model training device 900 also includes a fourth optimization module 950.

[0158] In some embodiments, the fourth optimization module 950 is configured to: decode multiple matching discrete vectors based on the decoder to obtain a first repaired image corresponding to the first image; optimize the decoder based on the first image and the first repaired image to obtain a target decoder; the second optimization module 930 is specifically configured to: determine the target image restoration model based on the target discrete codebook, the target prediction module and the target decoder.

[0159] like Fig. 9 As shown, the image restoration model training device 900 also includes a decoding module 960.

[0160] In some embodiments, the decoding module 960 is configured to: perform decoding processing on the multiple predicted discrete vectors based on the target decoder to obtain a second restored image corresponding to the second image.

[0161] In some embodiments, the decoding module 960 is specifically configured to: determine multiple enhanced predicted discrete vectors based on the second image features and multiple predicted discrete vectors; and decode the multiple enhanced predicted discrete vectors based on the target decoder to obtain a second restored image.

[0162] Fig.10 Schematic diagram of an image restoration device provided in some embodiments of the present application. Fig.10 As shown, the image restoration apparatus 1000 includes an encoding module 1010 , a prediction module 1020 and a decoding module 1030 .

[0163] The encoding module 1010 is configured to perform encoding processing on the target image based on the target encoder in the target image restoration model to obtain target image features corresponding to the target image.

[0164] The prediction module 1020 is configured to determine a plurality of target prediction discrete vectors corresponding to the target image features in the target discrete codebook based on the target prediction module in the target image restoration model.

[0165] The target discrete codebook includes multiple target discrete vectors, and the target discrete vectors include target prediction discrete vectors.

[0166] The decoding module 1030 is configured to decode the multiple target prediction discrete vectors based on the target decoder in the target image restoration model to obtain a target restoration image corresponding to the target image.

[0167] In some embodiments, the prediction module 1020 is specifically configured as: based on the target prediction module, determining multiple target prediction vector indices corresponding to the target image features in the target discrete codebook; wherein the target discrete codebook includes a target vector index corresponding to each target discrete vector in a plurality of target discrete vectors, and the target vector index includes a target prediction vector index; based on the plurality of target prediction vector indices, determining multiple target prediction discrete vectors corresponding to the target image features in the target discrete codebook.

[0168] In some embodiments, the decoding module 1030 is specifically configured to: determine multiple enhanced target prediction discrete vectors based on target image features and multiple target prediction discrete vectors; and decode the multiple enhanced target prediction discrete vectors based on the target decoder to obtain a target repaired image.

[0169] Figure 8 A schematic diagram of an electronic device provided for some embodiments of the present application. In some embodiments, the electronic device includes one or more processors and a memory. The memory is configured to store one or more programs. When the one or more programs are executed by the one or more processors, the one or more processors implement the training method of the image restoration model in the above-mentioned embodiment.

[0170] like Fig.11 As shown, the electronic device 1100 includes: a processor (processor) 1101 and a memory 1102. Exemplarily, the electronic device 1100 may also include: a communication interface (Communications Interface) 1103 and a communication bus 1104.

[0171] The processor 1101, the memory 1102 and the communication interface 1103 communicate with each other via the communication bus 1104. The communication interface 1103 is used to communicate with other devices such as client or other server network elements.

[0172] In some embodiments, the processor 1101 is used to execute the program 1105, and specifically can execute the relevant steps in the above-mentioned training method embodiment of the image restoration model. Specifically, the program 1105 may include a program code, and the program code includes computer executable instructions.

[0173] Exemplarily, the processor 1101 may be a central processing unit (CPU), or an application-specific integrated circuit (ASIC), or one or more integrated circuits configured to implement some embodiments of the present application. The electronic device 1100 may include one or more processors, which may be processors of the same type, such as one or more CPUs; or processors of different types, such as one or more CPUs and one or more ASICs.

[0174] In some embodiments, the memory 1102 is used to store the program 1105. The memory 1102 may include a high-speed RAM memory, and may also include a non-volatile memory (NVM), such as at least one disk memory.

[0175] Program 1105 can be specifically called by processor 1101 to enable electronic device 1100 to execute the training method operation of the image restoration model.

[0176] Some embodiments of the present application provide a computer-readable storage medium, which stores at least one executable instruction. When the executable instruction is executed on the electronic device 1100, the electronic device 1100 executes the training method of the image restoration model in the above embodiment.

[0177] The executable instructions can specifically be used to enable the electronic device 1100 to execute the training method operation of the image restoration model.

[0178] For example, the computer readable storage medium may be a read-only memory (ROM), a random access memory (RAM), a compact disc (CD-ROM), a magnetic tape, a floppy disk, an optical data storage device, and the like.

[0179] The beneficial effects that can be achieved by the readable storage medium provided in some embodiments of the present application can be referred to the beneficial effects in the training method of the corresponding image restoration model provided above, and will not be repeated here.

[0180] It should be noted that, in the application, relational terms such as first and second, etc. are only used to distinguish one entity or operation from another entity or operation, and do not necessarily require or imply any such actual relationship or order between these entities or operations. Moreover, the terms "include", "comprise" or any other variants thereof are intended to cover non-exclusive inclusion, so that a process, method, article or device including a series of elements includes not only those elements, but also other elements not explicitly listed, or also includes elements inherent to such process, method, article or device. In the absence of further restrictions, the elements defined by the sentence "including a ..." do not exclude the existence of other identical elements in the process, method, article or device including the elements.

[0181] Each embodiment in this specification is described in a related manner, and the same or similar parts between the embodiments can be referred to each other, and each embodiment focuses on the differences from other embodiments. In particular, for the device embodiment, since it is basically similar to the method embodiment, the description is relatively simple, and the relevant parts can be referred to the partial description of the method embodiment.

[0182] The logic and / or steps represented in the flowchart or otherwise described herein, for example, can be considered as an ordered list of executable instructions for implementing logical functions, which can be embodied in any computer-readable medium for use by an instruction execution system, apparatus or device (such as a computer-based system, a system including a processor or other system that can fetch instructions from an instruction execution system, apparatus or device and execute instructions), or used in combination with these instruction execution systems, apparatuses or devices.

[0183] For the purposes of this specification, a "computer-readable medium" can be any apparatus that can contain, store, communicate, propagate, or transport the program for use by an instruction execution system, apparatus, or device or in conjunction with such instruction execution system, apparatus, or device.

[0184] More specific examples (a non-exhaustive list) of computer-readable media include the following: an electrical connection having one or more wires (electronic device), a portable computer disk case (magnetic device), a random access memory (RAM), a read-only memory (ROM), an erasable and programmable read-only memory (EPROM or flash memory), a fiber optic device, and a portable compact disk read-only memory (CDROM).

[0185] In addition, the computer readable medium can even be a paper or other suitable medium on which the program can be printed, because the program can be obtained electronically, for example, by optically scanning the paper or other medium, and then editing, interpreting or processing in other suitable ways as necessary, and then storing it in the computer memory. It should be understood that various parts of the present application can be implemented in hardware, software, firmware or a combination thereof.

[0186] In the above embodiments, multiple steps or methods may be implemented by software or firmware stored in a memory and executed by a suitable instruction execution system. For example, if implemented by hardware, as in another embodiment, it may be implemented by any one of the following technologies known in the art or a combination thereof: a discrete logic circuit having a logic gate circuit for implementing a logic function on a data signal, a dedicated integrated circuit having a suitable combination of logic gate circuits, a programmable gate array (PGA), a field programmable gate array (FPGA), etc.

[0187] The above-described embodiments of the present application do not constitute a limitation on the protection scope of the present application.

Claims

1. A training method for an image restoration model, characterized in that: include: Determine a first image feature corresponding to the first image based on the image restoration model, and optimize multiple preset discrete vectors in a discrete codebook based on the first image feature to obtain a target discrete codebook; wherein the target discrete codebook includes multiple target discrete vectors; Determine a second image feature corresponding to the second image based on the image restoration model, and determine a plurality of predicted discrete vectors corresponding to the second image feature in the target discrete codebook based on a prediction module in the image restoration model; wherein the quality parameter of the first image is higher than the quality parameter of the second image, and the target discrete vector includes the predicted discrete vector; The prediction module is optimized based on the multiple target discrete vectors and the multiple predicted discrete vectors to obtain a target prediction module, and a trained target image restoration model is determined based on the target discrete codebook and the target prediction module.

2. The method according to claim 1, characterized in that: The optimizing a plurality of preset discrete vectors in a discrete codebook based on the first image feature to obtain a target discrete codebook includes: Determining a plurality of matching discrete vectors corresponding to the first image feature in the discrete codebook; wherein the preset discrete vectors include the matching discrete vectors; Based on the first image feature and the multiple matching discrete vectors, the multiple preset discrete vectors in the discrete codebook are optimized to obtain the target discrete codebook.

3. The method according to claim 2, characterized in that The determining, based on the image restoration model, a first image feature corresponding to the first image includes: Performing encoding processing on the first image based on the first encoder in the image restoration model to obtain the first image feature and a plurality of first image sub-features corresponding to the first image feature; The determining, in the discrete codebook, a plurality of matching discrete vectors corresponding to the first image feature comprises: Determining a candidate similarity between the first image sub-feature and each of the plurality of preset discrete vectors; Determine a target similarity with the highest similarity among the multiple candidate similarities, and determine a target preset discrete vector corresponding to the target similarity as the matching discrete vector corresponding to the first image sub-feature; wherein the preset discrete vector includes the target preset discrete vector.

4. The method according to claim 1, characterized in that The determining, based on the prediction module in the image restoration model, a plurality of predicted discrete vectors corresponding to the second image feature in the target discrete codebook comprises: Determining, based on the prediction module, a plurality of prediction vector indexes corresponding to the second image feature in the target discrete codebook; wherein the target discrete codebook includes target vector indexes corresponding to each of the target discrete vectors, and the target vector indexes include the prediction vector indexes; The plurality of prediction discrete vectors corresponding to the second image feature are determined in the target discrete codebook based on the plurality of prediction vector indices.

5. The method according to claim 4, characterized in that The determining, based on the prediction module, a plurality of prediction vector indexes corresponding to the second image feature in the target discrete codebook comprises: Determining a high-level feature corresponding to the second image feature based on the prediction module; wherein the high-level feature is determined by a self-attention mechanism; The plurality of prediction vector indexes corresponding to the high-level features are determined in the target discrete codebook based on the prediction module.

6. The method according to claim 4, characterized in that The step of optimizing the prediction module based on the multiple target discrete vectors and the multiple prediction discrete vectors to obtain a target prediction module includes: Determine a code-level loss function based on the multiple target vector indices corresponding to the multiple target discrete vectors and the multiple prediction vector indices corresponding to the multiple prediction discrete vectors; The prediction module is optimized based on the code-level loss function to obtain the target prediction module.

7. The method according to any one of claims 1 to 6, characterized in that The determining, based on the image restoration model, a second image feature corresponding to the second image includes: The second image is encoded based on the second encoder in the image restoration model to obtain the second image features corresponding to the second image.

8. The method according to claim 7, characterized in that The method further comprises: determining a feature-level loss function based on the plurality of target discrete vectors and the second image features; Optimizing the second encoder based on the feature-level loss function to obtain a target encoder; The step of determining a trained target image restoration model based on the target discrete codebook and the target prediction module includes: The target image restoration model is determined based on the target discrete codebook, the target prediction module and the target encoder.

9. The method according to any one of claims 1 to 6, characterized in that: The method further comprises: Decoding the plurality of matching discrete vectors based on a decoder to obtain a first restored image corresponding to the first image; Optimizing the decoder based on the first image and the first repaired image to obtain a target decoder; The step of determining a trained target image restoration model based on the target discrete codebook and the target prediction module includes: The target image restoration model is determined based on the target discrete codebook, the target prediction module and the target decoder.

10. The method according to claim 9, characterized in that The method further comprises: The plurality of predicted discrete vectors are decoded based on the target decoder to obtain a second restored image corresponding to the second image.

11. The method according to claim 10, characterized in that The decoding process of the plurality of predicted discrete vectors based on the target decoder to obtain a second restored image corresponding to the second image includes: Determining a plurality of enhanced predicted discrete vectors based on the second image feature and the plurality of predicted discrete vectors; The plurality of enhanced predicted discrete vectors are decoded based on the target decoder to obtain the second restored image.

12. An image restoration method, characterized in that: include: Encoding the target image based on the target encoder in the target image restoration model to obtain target image features corresponding to the target image; Based on the target prediction module in the target image restoration model, a plurality of target prediction discrete vectors corresponding to the target image features are determined in a target discrete codebook; wherein the target discrete codebook includes a plurality of target discrete vectors, and the target discrete vectors include the target prediction discrete vector; Based on the target decoder in the target image restoration model, the multiple target prediction discrete vectors are decoded to obtain a target restoration image corresponding to the target image.

13. The method according to claim 12, characterized in that The target prediction module in the target image restoration model determines a plurality of target prediction discrete vectors corresponding to the target image features in a target discrete codebook, including: Based on the target prediction module, a plurality of target prediction vector indexes corresponding to the target image features are determined in the target discrete codebook; wherein the target discrete codebook includes a target vector index corresponding to each of the target discrete vectors in the plurality of target discrete vectors, and the target vector index includes the target prediction vector index; The multiple target prediction discrete vectors corresponding to the target image features are determined in the target discrete codebook based on the multiple target prediction vector indices.

14. The method according to claim 12, characterized in that The target decoder in the target image restoration model decodes the multiple target prediction discrete vectors to obtain a target restoration image corresponding to the target image, including: Determining a plurality of enhanced target prediction discrete vectors based on the target image feature and the plurality of target prediction discrete vectors; The target decoder decodes the multiple enhanced target prediction discrete vectors to obtain the target repaired image.

15. A training device for an image restoration model, characterized in that: include: A first optimization module is configured to determine a first image feature corresponding to the first image based on the image restoration model, and optimize a plurality of preset discrete vectors in a discrete codebook based on the first image feature to obtain a target discrete codebook; wherein the target discrete codebook includes a plurality of target discrete vectors; a determination module configured to determine a second image feature corresponding to the second image based on the image restoration model, and determine a plurality of predicted discrete vectors corresponding to the second image feature in the target discrete codebook based on a prediction module in the image restoration model; wherein the quality parameter of the first image is higher than the quality parameter of the second image, and the target discrete vector includes the predicted discrete vector; The second optimization module is configured to optimize the prediction module based on the multiple target discrete vectors and the multiple predicted discrete vectors to obtain a target prediction module, and determine a trained target image restoration model based on the target discrete codebook and the target prediction module.

16. An image restoration device, characterized in that: include: An encoding module is configured to perform encoding processing on a target image based on a target encoder in a target image restoration model to obtain a target image feature corresponding to the target image; A prediction module is configured to determine, based on a target prediction module in the target image restoration model, a plurality of target prediction discrete vectors corresponding to the target image features in a target discrete codebook; wherein the target discrete codebook includes a plurality of target discrete vectors, and the target discrete vectors include the target prediction discrete vector; The decoding module is configured to decode the multiple target prediction discrete vectors based on the target decoder in the target image restoration model to obtain a target restoration image corresponding to the target image.

17. An electronic device, characterized in that: include: one or more processors; and A memory configured to: store one or more programs; Wherein, when the one or more programs are executed by the one or more processors, the one or more processors implement the training method of the image restoration model according to any one of claims 1-11, or the image restoration method according to any one of claims 12-14.

18. A computer-readable storage medium, characterized in that: A computer program is stored thereon, and when the computer program is executed by a processor, the training method of the image restoration model according to any one of claims 1-11, or the image restoration method according to any one of claims 12-14 is implemented.

Images