Image restoration model training method, device, computer equipment and storage medium

Abstract

The present application discloses a training method, device, computer equipment and storage medium for an image restoration model, the training method for the image restoration model comprising: obtaining a mask image corresponding to a sample image; extracting an edge image of the sample image as a first edge image; performing edge restoration on the first edge image based on a pre-trained edge restoration model, and obtaining the restored edge image as a second edge image; obtaining a restoration mask image and a context mask image corresponding to the sample image based on the mask image, the restoration mask image being used to indicate an area to be restored in the sample image, and the context mask image being used to indicate a known image area when restoring the area to be restored; training an initial model based on the sample image, the restoration mask image, the context mask image and the second edge image, and obtaining a trained image restoration model. The method can improve the image restoration effect of the image restoration model.

Description

Technical Field

[0001] The present application relates to the field of image processing technology, and more specifically, to a training method, apparatus, computer equipment and storage medium for an image restoration model. Background Art

[0002] With the rapid advancement of science and technology and living standards, electronic devices (such as smart phones, tablets, etc.) are widely used by people, and the shooting functions of electronic devices are becoming more and more powerful. Non-professionals can also use electronic devices to shoot high-quality images. However, during the process of image shooting, transmission, processing and storage, due to various factors, the image may be missing, cracked, etc. Therefore, methods for repairing these defective images have emerged, but the image repair effect of current image repair methods is not good. Summary of the invention

[0003] In view of the above problems, the present application proposes a training method, apparatus, computer device and storage medium for an image restoration model.

[0004] In a first aspect, an embodiment of the present application provides a method for training an image restoration model, the method comprising: obtaining a mask image corresponding to a sample image, the mask image being used to indicate an area where a significant object in the sample image is located; extracting an edge image of the sample image as a first edge image, the edge image being used to characterize the texture of the image content of the sample image; performing edge completion on the first edge image based on a pre-trained edge completion model to obtain the completed edge image as a second edge image; and training an initial model based on the sample image, the mask image and the second edge image to obtain a trained image restoration model.

[0005] In a second aspect, an embodiment of the present application provides a training device for an image restoration model, the device comprising: a mask acquisition module, a first edge acquisition module, a second edge acquisition module and a model training module, wherein the mask acquisition module is used to acquire a mask image corresponding to a sample image, the mask image being used to indicate an area where a significant object in the sample image is located; the first edge acquisition module is used to extract an edge image of the sample image as a first edge image, the edge image being used to characterize a texture of the image content of the sample image; the second edge acquisition module is used to perform edge completion on the first edge image based on a pre-trained edge completion model, and obtain the completed edge image as a second edge image; the model training module is used to train an initial model based on the sample image, the mask image and the second edge image, and obtain a trained image restoration model.

[0006] In a third aspect, an embodiment of the present application provides a computer device comprising: one or more processors; a memory; one or more applications, wherein the one or more applications are stored in the memory and configured to be executed by the one or more processors, and the one or more programs are configured to execute the training method of the image restoration model provided in the first aspect above.

[0007] In a fourth aspect, an embodiment of the present application provides a computer-readable storage medium, in which a program code is stored. The program code can be called by a processor to execute the training method of the image restoration model provided in the first aspect above.

[0008] The solution provided by the present application is to obtain a mask image corresponding to a sample image, the mask image is used to indicate the area where a significant object is located in the sample image, extract the edge image of the sample image as the first edge image, the edge image is used to characterize the texture of the image content of the sample image, perform edge completion on the first edge image based on a pre-trained edge completion model, obtain the completed edge image as the second edge image, train the initial model based on the sample image, the mask image and the second edge image, and obtain a trained image restoration model. Since the edge completion is performed on the extracted edge image during model training, and the area to be restored and the known image area are refined based on the mask image, the trained model can better restore the image content and improve the image restoration effect. BRIEF DESCRIPTION OF THE DRAWINGS

[0009] In order to more clearly illustrate the technical solutions in the embodiments of the present application, the drawings required for use in the description of the embodiments will be briefly introduced below. Obviously, the drawings described below are only some embodiments of the present application. For those skilled in the art, other drawings can be obtained based on these drawings without creative work.

[0010] Figure 1 A flow chart of a method for training an image restoration model according to an embodiment of the present application is shown.

[0011] Figure 2 A schematic diagram of a mask image of a sample image provided in an embodiment of the present application is shown.

[0012] Figure 3 A schematic diagram of the structure of an edge completion model provided in an embodiment of the present application is shown.

[0013] Figure 4 A flow chart of a method for training an image restoration model according to another embodiment of the present application is shown.

[0014] Figure 5A flowchart of step S260 in a method for training an image restoration model provided in another embodiment of the present application is shown.

[0015] Figure 6 A schematic diagram of an image during the model training process provided in an embodiment of the present application is shown.

[0016] Figure 7 A schematic diagram of the structure of the initial model provided in an embodiment of the present application is shown.

[0017] Figure 8 A schematic diagram of the structure of the middle layer in the initial model provided in an embodiment of the present application is shown.

[0018] Fig. 9 A flow chart of a method for training an image restoration model according to yet another embodiment of the present application is shown.

[0019] Fig.10 A schematic diagram showing the effect of processing a mask image provided by an embodiment of the present application is shown.

[0020] Fig.11 A flow chart of a method for training an image restoration model according to yet another embodiment of the present application is shown.

[0021] Fig.12 A schematic diagram showing the effect of processing a captured image provided by an embodiment of the present application is shown.

[0022] Fig.13 A block diagram of a training device for an image restoration model according to an embodiment of the present application is shown.

[0023] Fig.14 It is a block diagram of a computer device according to an embodiment of the present application for executing a training method for an image restoration model according to an embodiment of the present application.

[0024] Fig.15 It is a storage unit of an embodiment of the present application for storing or carrying program code for implementing a training method for an image restoration model according to an embodiment of the present application. DETAILED DESCRIPTION

[0025] In order to enable those skilled in the art to better understand the solution of the present application, the technical solution in the embodiments of the present application will be clearly and completely described below in conjunction with the drawings in the embodiments of the present application.

[0026] Image restoration is a very important research topic in image processing. Image restoration refers to restoring a damaged image to a state similar to the original state through certain technical means. At present, image restoration methods have been widely used in the fields of old picture restoration, photo beautification, background character removal, 3D photography, etc.

[0027] In the related art, the image restoration methods include traditional image restoration algorithms and deep learning-based image restoration methods. Among them, the traditional image restoration algorithm uses algorithms such as dilation to fill in the area to be restored using pixels from the adjacent area; the image restoration method based on deep learning uses convolutional neural networks to learn image restoration capabilities, uses the network to obtain image semantic information and fills the area to be restored.

[0028] However, traditional image restoration algorithms rely on neighboring pixels, and the restoration results are generally blurry and have no clear texture. Current deep learning-based image restoration algorithms are prone to errors in some common sense areas or complex areas. Therefore, the restoration effect of current image restoration algorithms is still poor.

[0029] In response to the above problems, the inventors have proposed a training method, device, computer equipment, and storage medium for an image restoration model provided in the embodiments of the present application, which can complete the edges of the extracted edge image during model training, and refine the area to be restored and the known image area based on the mask image, so that the trained model can better restore the image content and improve the image restoration effect. The specific training method of the image restoration model will be described in detail in the subsequent embodiments.

[0030] See also Figure 1 , Figure 1 FIG. 1 is a flow chart of a training method for an image restoration model provided by an embodiment of the present application. In a specific embodiment, the training method for the image restoration model is applied to Fig.13 The training device 400 of the image restoration model and the computer device 100 ( Fig.14 ). The following will take a computer device as an example to illustrate the specific process of this embodiment. Of course, it can be understood that the computer device used in this embodiment can be a server, a laptop, a PC, a smart phone, a tablet computer, etc., which is not limited here. Figure 1 The process shown is described in detail, and the training method of the image restoration model may specifically include the following steps:

[0031] Step S110: obtaining a mask image corresponding to the sample image, wherein the mask image is used to indicate an area where a significant object in the sample image is located.

[0032] In an embodiment of the present application, a computer device may detect significant objects on a sample image based on a sample image containing significant objects, thereby obtaining a mask image corresponding to the sample image. The significant object refers to an object region in an image that can attract the visual attention of a user, and significant object detection is to detect the object region in an image that can attract the visual attention of a user. The mask image is a binary mask image, which refers to a binary image obtained by segmenting the sample image according to the pixel values ​​of the image. In the mask image, the pixel values ​​of the region where the significant object is located are different from those of other regions, so that the region where the significant object is located in the mask image can be distinguished. For example, the mask image corresponding to the sample image is as follows: Figure 2 As shown, the pixel value corresponding to the area where the salient object is located can be 255 (white), while the pixel value of other areas can be 0 (black). It can be understood that since image restoration is usually performed on the edge or nearby area of ​​some target object area, for example, in 3D photography, the area on the edge of the portrait is usually restored, so based on the sample image containing the salient object, the corresponding mask image can be obtained, and the model can be trained based on the mask image to improve the restoration effect of the image restoration model.

[0033] In some embodiments, the computer device may obtain an image sample set, which includes a plurality of sample images containing salient objects and having complete image content. The sample images may be obtained from an open source image library to construct the image sample set. Of course, the specific method of obtaining the sample images may not be limited. For example, the sample images may be obtained by photographing the salient objects. After obtaining the image sample set, the computer device may obtain the mask images corresponding to the sample images in the image sample set.

[0034] As a possible implementation, the sample image may be input into a pre-trained salient object detection model. The salient object detection model is pre-selected and trained to detect the region where the salient object is located based on the input image, and then output a mask image corresponding to the region where the salient object is located. Thus, after the sample image is input into the pre-trained salient object detection model, a mask image corresponding to the region where the salient object is located output by the salient object detection model may be obtained.

[0035] Optionally, the salient object detection model for the above-mentioned mask image may include an encoding module and a decoding module. That is, the salient object detection model may be designed as an Encoder-Decoder structure, where the Encoder stage continuously downsamples the image to extract deep feature information of the image, the Decoder stage upsamples the image, and finally performs a binary classification process, and finally outputs a detection result of a salient object with the same size as the input (resolution is also the second resolution). The detection result of the salient object may include a mask image corresponding to the area where the salient object is located.

[0036] Optionally, the sample image may be preprocessed before being input into the salient object detection model. Optionally, the salient object detection model requires that the pixel values ​​of the input image should be normalized, for example, the pixel values ​​should be real numbers between [0, 1]. If the pixel values ​​of the resized image are not normalized, they should be normalized. For example, the pixel values ​​of the resized image are represented as integers between [0, 255], which can be normalized by dividing by 255 so that the salient object detection model can process it correctly. It is understandable that normalization can have different definitions. For example, in another normalization definition, the pixel value should be a real number between [-1, 1]. For different normalization definitions, the normalization method should be adjusted accordingly.

[0037] In some cases, preprocessing the sample image may further include: subtracting the average pixel value from the pixel value of the sample image, and dividing the pixel value by the pixel value variance. It can be understood that in order to standardize the sample image, the average pixel value of the pixel point of the sample image may be subtracted from the pixel value of each pixel point of the sample image. In addition, the variance of the pixel value of the pixel point may be calculated to obtain the pixel value variance, and after the pixel value of each pixel point is subtracted from the average pixel value of the pixel point of the sample image, the pixel value of each pixel point may be divided by the pixel value variance.

[0038] In the above implementation, when training the salient object detection model, the untrained model can be trained based on the image containing the salient object and the result of labeling the salient object in the image, and the loss value is calculated by the result output by the untrained model and the result after labeling the shadows in the image, and then the parameters of the untrained model are adjusted according to the loss value until the loss value meets the conditions, so that the salient object detection model can correctly detect the area where the salient object is located.

[0039] As another possible implementation, the image restoration model that needs to be trained in the embodiment of the present application is for restoration of 3D photography, and the above sample images are sample images annotated with saliency detection results in the public data set MS-COCO, wherein the saliency detection results may include the area where the salient objects are located. Therefore, when acquiring the above sample images, the saliency detection results corresponding to the annotations of the sample images can be acquired at the same time, and based on the saliency detection results corresponding to the annotations of the sample images, a mask image corresponding to the sample images can be generated.

[0040] Of course, the specific manner in which the computer device obtains the mask image corresponding to the sample image may not be limited.

[0041] Step S120: extracting an edge image of the sample image as a first edge image, wherein the edge image is used to characterize the texture of the image content of the sample image.

[0042] In an embodiment of the present application, the computer device can extract the edge image of the sample image to obtain the texture features in the sample image, so that when repairing the image, the missing part can be repaired based on the known edges.

[0043] In some embodiments, the edge image may be obtained by extracting the texture of the entire sample image. The edge information of the sample image may be extracted using an edge detection algorithm to obtain the edge image. Optionally, the edge detection algorithm may be a Canny edge detection operator, a Sobel edge detection operator, etc., which are not limited here.

[0044] In some embodiments, the edge image can also be obtained by a pre-trained edge extraction model. The sample image can be input into the pre-trained edge extraction model, and then the edge image output by the edge extraction model is obtained and used as the first edge image. The edge extraction model can be an encoder-decoder model, an hourglass network, an autoencoder network, etc., and the specific type of the model may not be limited.

[0045] Step S130: Based on a pre-trained edge completion model, edge completion is performed on the first edge image to obtain a completed edge image as a second edge image.

[0046] In an embodiment of the present application, after obtaining the first edge image of the sample image, edge completion can be performed on the first edge image based on a pre-trained edge completion model to achieve image texture repair. Therefore, when performing subsequent image repair, texture loss in the repaired image, such as missing straight lines, can be avoided.

[0047] In some embodiments, when edge completion is performed on the first edge image based on a pre-trained edge completion model, the first edge image, the sample image, and the mask image can be input into the pre-trained edge completion model to obtain an image output by the edge completion model as the second edge image. The edge extraction model can be an encoder-decoder model, an hourglass network, an autoencoder network, etc., and the specific type of the model may not be limited.

[0048] As a possible implementation, the edge completion model may be an encoder-decoder model. In the encoder-decoder model, the encoder may be used to extract image features of an input image, and the decoder may be used to output a new image based on the input image features. As an example, the edge completion model may be as follows: Figure 3 As shown in the figure, it includes an encoder, an intermediate layer, and a decoder. The encoder first downsamples the input image three times to extract image features. The intermediate layer uses 8 ResNet modules. The unique residual structure of the ResNet module can simultaneously learn the original image features and the features after convolution, greatly increasing the accumulation of features. Finally, the decoder uses the opposite structure of the encoder to upsample to restore the original image size and obtain a new image (the completed edge image).

[0049] Optionally, when training the edge completion model, a discriminator D1 can be added. During the training process, the completed edge image can be input into the discriminator D1 to determine whether it is a real image, so that the output value of the discriminator can be used to constrain the model training according to the output of the discriminator, that is, the encoder-decoder model and the discriminator are allowed to play a game. Finally, the edge completion model obtained after the training is completed can output a more realistic image. It should be noted that before the first edge image, the sample image and the mask image are input into the pre-trained edge completion model, the sample image needs to be converted into a grayscale image, and then input into the model together with the first edge image and the mask image.

[0050] Of course, the specific method of performing edge completion on the first edge image may not be limited.

[0051] Step S140: Based on the mask image, obtain a repair mask image and a context mask image corresponding to the sample image, wherein the repair mask image is used to indicate an area to be repaired in the sample image, and the context mask image is used to indicate a known image area when repairing the area to be repaired.

[0052] In an embodiment of the present application, the computer device can also obtain a repair mask image and a context mask image corresponding to the sample image based on the mask image, thereby obtaining the area to be repaired required when repairing the image, and the known image area used to repair the area to be repaired.

[0053] In some embodiments, the computer device may determine the edge of the area where the significant object is located, the area adjacent to the edge in the area, or the area adjacent to the edge in other areas as the area to be repaired, based on the area where the significant object is located indicated in the mask image, and then obtain the repair mask image. For the context mask image, based on the area to be repaired indicated by the repair mask image, the area adjacent to the area to be repaired may be determined as the known image area when repairing the area to be repaired. Of course, the specific method of obtaining the repair mask image and the context mask image corresponding to the sample image in the embodiment of the present application may not be limited. For example, it may also be implemented using the implementation method in the subsequent embodiments.

[0054] Step S150: Based on the sample image, the restoration mask image, the context mask image and the second edge image, the initial model is trained to obtain a trained image restoration model.

[0055] In an embodiment of the present application, after obtaining the repair mask image, the context mask image and the second edge image, the initial model can be trained based on the sample image, the repair mask image, the context mask image and the second edge image to obtain a trained image repair model. The initial model can be an encoder-decoder model, an hourglass network, an autoencoder network, a generative adversarial model, etc., which are not limited here.

[0056] In some embodiments, since it is necessary to know the area to be repaired and the content of the known image area during image repair, the content of the known image area can be determined based on the sample image and the context mask image, and then integrated with the repair mask image into an image, so that the obtained image contains both the content of the known image area used to repair the area to be repaired and the image area to be repaired; then the obtained image and the second edge image are input into the initial model to obtain the output result of the initial model. Since the sample image is an image with complete image content, that is, the area to be repaired in the sample image is actually complete, the content of the area to be repaired in the output result should be similar to the content of the area to be repaired in the sample image. Based on this, the training of the initial model can be constrained, and then an image repair model that can accurately repair the image can be trained.

[0057] The training method of the image restoration model provided in the embodiment of the present application is that during model training, the edge image of the extracted sample image is edge-filled, and the area to be restored and the known image area are determined based on the mask image of the sample image, and then input into the model for training. Therefore, the trained model can better restore the image content and improve the image restoration effect.

[0058] See also Figure 4 , Figure 4 A flow chart of a training method for an image restoration model provided by another embodiment of the present application is shown. The training method for the image restoration model is applied to the above-mentioned computer device. Figure 4 The process shown is described in detail, and the training method of the image restoration model may specifically include the following steps:

[0059] Step S210: obtaining a mask image corresponding to the sample image, wherein the mask image is used to indicate an area where a significant object in the sample image is located.

[0060] Step S220: extracting an edge image of the sample image as a first edge image, wherein the edge image is used to characterize the texture of the image content of the sample image.

[0061] Step S230: Based on the pre-trained edge completion model, edge completion is performed on the first edge image to obtain a completed edge image as a second edge image.

[0062] Step S240: Based on the mask image, obtain a repair mask image and a context mask image corresponding to the sample image, wherein the repair mask image is used to indicate an area to be repaired in the sample image, and the context mask image is used to indicate a known image area when repairing the area to be repaired.

[0063] In the embodiment of the present application, steps S210 to S240 may refer to the contents of other embodiments and will not be described in detail here.

[0064] Step S250: Based on the repair mask image and the context mask image, a background area image is obtained, where the background area image is used to indicate the area to be repaired and the known image area.

[0065] In an embodiment of the present application, after obtaining the repair mask image, the context mask image and the second edge image, the repair mask image and the context mask image may be processed before training the initial model based on the sample image, the repair mask image, the context mask image and the second edge image, so as to train a more accurate image repair model. Among them, a background area image may be obtained based on the repair mask image and the context mask image, and the background area image is used to indicate the area to be repaired and the known image area, so that the loss value in the model training can be accurately calculated based on the background area image later.

[0066] In some embodiments, the computer device may add the repair mask image to the context mask image to obtain a background area image. It can be understood that the area to be repaired indicated in the repair mask image and the known image area indicated in the context mask image are adjacent to each other, so by adding the two, a mask image including both the area to be repaired and the known image area can be obtained, and the mask image is used as the background area image.

[0067] Step S260: Based on the sample image, the restoration mask image, the context mask image, the background area image and the second edge image, the initial model is trained to obtain a trained image restoration model.

[0068] In an embodiment of the present application, after obtaining the background area image, the initial model can be trained based on the sample image, the repair mask image, the context mask image, the background area image and the second edge image to obtain a trained image repair model. The background area image can be used to determine the area for calculating the loss during the training process, and the area indicated in the background area image includes the area to be repaired and the known image area, so the area for calculating the loss can be refined without calculating the loss of other irrelevant areas, so that the loss can be accurately calculated, and the model can be trained based on this to obtain a more accurate image repair model.

[0069] In some embodiments, please participate Figure 5 , based on the sample image, the repair mask image, the context mask image, the background area image and the second edge image, the initial model is trained to obtain a trained image repair model, which may include steps S261 to S266.

[0070] Step S261: multiply the background area image and the sample image to obtain a first image.

[0071] In this embodiment, the background area image and the sample image may be multiplied to obtain a first image, whereby the obtained first image includes the image content of the original area to be restored in the sample image and the image content of the known image area.

[0072] Step S262: Add the first image and the repair mask image to obtain a second image.

[0073] In this embodiment, after obtaining the first image, the first image is added to the repair mask image to obtain the second image. For the area to be repaired, the repair mask image is a white area, and the first image is the actual image content of the sample image. After adding the first image to the repair mask image, the image content corresponding to the area to be repaired in the first image can be changed to a white area, so that the second image contains the image content of the known image area in the sample image, and the area to be repaired (white area) adjacent to the known image area, while other areas are black areas. It can be understood that when the first image is added to the repair mask image, since the area to be repaired in the repair mask image is white, that is, the grayscale value is 255, when it is added to the grayscale value of any pixel, the resulting grayscale value is greater than 255, and when the images are added, if the grayscale value after the addition of the pixels is greater than 255, 255 will be taken. Therefore, after adding the first image to the repair mask image, the second image can contain the image content of the above-mentioned known image area in the sample image, as well as the area to be repaired (white area) adjacent to the known image area, while other areas are black areas.

[0074] Step S263: Connect the second image and the second edge image to obtain an input image.

[0075] Step S264: input the input image to the initial model to obtain an output image output by the initial model.

[0076] In this embodiment, after obtaining the second image, the second image and the second edge image can be connected to obtain an input image for inputting the initial model, and the input image is input into the initial model to obtain an output image output by the initial model. The obtained output image is equivalent to the image content after the area to be repaired (white area) in the second image is changed to the repaired area, that is, the output image contains the image content of the area to be repaired after the repair and the image content of the known image area, while other areas are black areas. Figure 6 As shown, it shows a schematic diagram of the change of the sample image, the first column of images is the sample image, the second column of images is the above-mentioned second image, the third column of images is the above-mentioned second edge image, and the fourth column of images is the output image.

[0077] Step S265: Based on the repair mask image and the background area image, a loss value of the output image relative to the sample image is obtained as a total loss value.

[0078] After obtaining the output image output by the initial model, the loss value of the output image relative to the sample image can be obtained as the total loss value based on the repaired mask image and the background area image, so as to iteratively train the initial model according to the total loss value.

[0079] Among them, the computer device can obtain the area indicated by the repair mask image from the sample image as the first area image, and obtain the area indicated by the repair mask image from the output image as the second area image; obtain the difference between the second area image and the first area image as the first loss value; obtain the area indicated by the background area image from the sample image as the third area image, and obtain the area indicated by the background area image from the output image as the fourth area image; obtain the difference between the fourth area image and the third area image as the second loss value; based on the first loss value and the second loss value, obtain the loss value of the output image relative to the sample image as the total loss value. Optionally, the first loss value and the second loss value can be weighted and summed to obtain the total loss value. Thus, it can be achieved that the two parts of loss are determined based on the changes of the integrated area of ​​the area to be repaired and the known image area after repair relative to the sample image before repair, and the changes of the area to be repaired after repair relative to the sample image before repair, and then the total loss value is determined, so that the total loss value can be determined more accurately, so that an accurate model can be trained.

[0080] Step S266: Based on the total loss value, the initial model is iteratively trained until the initial model meets preset conditions, thereby obtaining a trained image restoration model.

[0081] Optionally, the parameters in the initial model can be updated according to the gradient of the back-propagation loss of the back-propagation algorithm until the preset conditions are met. The preset conditions may be: the total loss value is less than the preset value, the total loss value no longer changes, or the number of training times reaches the preset number of times, etc. It can be understood that after the initial model is iteratively trained for multiple training cycles according to the image sample set, wherein each training cycle includes multiple iterative trainings, and the parameters of the initial model are continuously optimized, the above total loss value becomes smaller and smaller, and finally becomes a fixed value, or is less than the above preset value, at which point, it indicates that the initial model has converged; of course, it can also be determined that the initial model has converged after the number of training times reaches the preset number, thereby obtaining a trained image restoration model.

[0082] As a possible implementation, the above initial model may be an encoder-decoder model. In the initial model, the encoder may be used to extract image features of the input image, and the decoder may be used to output a new image based on the input image features. As an example, the edge completion model may be as follows: Figure 7 As shown in the figure, it includes an encoder, an intermediate layer and a decoder; the encoder first downsamples the input image three times to extract image features; the intermediate layer uses 8 ResNet modules, and uses the unique residual structure of the ResNet module to simultaneously learn the original image features and the features after convolution, greatly increasing the accumulation of features; finally, the decoder uses the opposite structure of the encoder to upsample to restore the original image size and obtain a new image (output image). The ResNet module can be as follows Figure 8 As shown in the figure, it can simultaneously learn the original image feature X and the convolution feature F(X) to obtain the accumulated feature F(X)+X.

[0083] Optionally, when training the above initial model, a discriminator D2 can be added. During the training process, the output image can be input into the discriminator D2 to determine whether it is a real image, so that the output value of the discriminator can be used to constrain the training of the initial model according to the output value of the discriminator, that is, let the encoder-decoder model and the discriminator play a game. Finally, the image restoration model obtained after the training is completed makes its output image more realistic.

[0084] As an example, the size of the second image is 128*128*1, and the size of the second edge image is 128*128*1. After splicing the two, an input image of 128*128*2 is obtained; the output feature map size of the first convolutional layer of the encoder becomes 64*64*64, the second convolutional layer uses the output feature map of the first convolutional layer as input, and the output feature map size is 32*32*128, and the third convolutional layer uses the output feature map of the second convolutional layer as input, and the output feature map size is 16*16*256; in the middle layer, the size of the input and output feature maps of each layer is 16*16*256; the first deconvolution layer of the decoder takes the feature map output from the middle layer as output, and outputs a feature map of 32*32*128; the second deconvolution layer of the decoder takes the output feature map of the first deconvolution layer as input, and outputs a feature map of 64*64*64; the decoder convolution layer restores the 64*64*64 feature map to an output image of 128*128*1, which is the repair result.

[0085] The training method of the image restoration model provided in the embodiment of the present application performs edge completion on the edge image of the extracted sample image during the model training process, and determines the area to be restored and the known image area based on the mask image of the sample image, and then inputs them into the model for training, so that the trained model can better restore the image content and improve the image restoration effect. Furthermore, the background area image is determined through the area to be restored and the known image area, and the loss value is determined based on the background area image and the area to be restored, so that the loss of the model can be accurately calculated, thereby improving the accuracy of the trained image model.

[0086] See also Fig. 9 , Fig. 9 A flow chart of a training method for an image restoration model provided by another embodiment of the present application is shown. The training method for the image restoration model is applied to the above-mentioned computer device. Fig. 9 The process shown is described in detail, and the training method of the image restoration model may specifically include the following steps:

[0087] Step S310: obtaining a mask image corresponding to a sample image, wherein the mask image is used to indicate an area where a significant object in the sample image is located.

[0088] Step S320: extracting an edge image of the sample image as a first edge image, wherein the edge image is used to characterize the texture of the image content of the sample image.

[0089] Step S330: Based on the pre-trained edge completion model, edge completion is performed on the first edge image to obtain a completed edge image as a second edge image.

[0090] In the embodiment of the present application, steps S310 to S330 refer to the contents of other embodiments and will not be repeated here.

[0091] Step S340: dilating the mask image to obtain a first processed image.

[0092] Step S350: performing corrosion processing on the mask image to obtain a second processed image.

[0093] In the embodiment of the present application, after obtaining the mask image corresponding to the sample image, the mask image can be dilated and eroded respectively, thereby obtaining a first processed image obtained by the dilation process and a second processed image obtained by the erosion process. It can be understood that by dilating and eroding the above-mentioned mask image, two mask images (i.e., the first processed image and the second processed image) with adjacent edges can be obtained, and the image dilation and erosion can maintain the approximate contour of the image, thereby obtaining a ring-shaped repair area similar to the shape of the original mask image, so that when the area to be repaired is determined based on the first processed image and the second processed image, the image repair area can be increased as much as possible.

[0094] Step S360: Based on the mask image, the first processed image and the second processed image, obtain a repair mask image and a context mask image corresponding to the sample image, wherein the repair mask image is used to indicate an area to be repaired in the sample image, and the context mask image is used to indicate a known image area when repairing the area to be repaired.

[0095] In an embodiment of the present application, after obtaining the above-mentioned mask image, the first processed image and the second processed image, the repaired mask image and the context mask image corresponding to the sample image can be obtained based on the above-mentioned mask image, the first processed image and the second processed image.

[0096] As a possible implementation, Fig.10 As shown, the first processed image and the mask image can be bitwise ANDed (XORed) to obtain a context mask image corresponding to the sample image; the second processed image and the mask image can be bitwise ANDed (XORed) to obtain a repaired mask image corresponding to the sample image.

[0097] As another possible implementation, the first processed image and the mask image may be bitwise ANDed to obtain a repaired mask image corresponding to the sample image; and the second processed image and the mask image may be bitwise ANDed to obtain a context mask image corresponding to the sample image.

[0098] Step S370: Based on the sample image, the restoration mask image, the context mask image and the second edge image, the initial model is trained to obtain a trained image restoration model.

[0099] In the embodiment of the present application, step S370 can refer to the content of the above embodiment, which will not be repeated here. If, according to the previous embodiment, the background area image is obtained based on the repair mask image and the context mask image, and the model training is performed based on the background area image, the background area image can be as follows: Fig.10 as shown in .

[0100] The training method of the image restoration model provided in the embodiment of the present application performs edge completion on the edge image of the extracted sample image during the model training process, and determines the area to be restored and the known image area based on the mask image of the sample image, and then inputs them into the model for training, so that the trained model can better restore the image content and improve the image restoration effect. Furthermore, by dilating and corroding the mask image, and obtaining the restoration mask image and the context mask image based on the processed image, the area to be restored can be enlarged, and the restoration effect of the image restoration model can be improved.

[0101] See also Fig.11 , Fig.11 A flow chart of a training method for an image restoration model provided by another embodiment of the present application is shown. The training method for the image restoration model is applied to the above-mentioned computer device. Fig.11 The process shown is described in detail, and the training method of the image restoration model may specifically include the following steps:

[0102] Step S401: obtaining a mask image corresponding to a sample image, wherein the mask image is used to indicate an area where a significant object in the sample image is located;

[0103] Step S402: extracting an edge image of the sample image as a first edge image, wherein the edge image is used to characterize the texture of the image content of the sample image;

[0104] Step S403: performing edge completion on the first edge image based on a pre-trained edge completion model to obtain a completed edge image as a second edge image;

[0105] Step S404: Based on the mask image, a repair mask image and a context mask image corresponding to the sample image are obtained, wherein the repair mask image is used to indicate an area to be repaired in the sample image, and the context mask image is used to indicate a known image area when repairing the area to be repaired;

[0106] Step S405: Based on the sample image, the restoration mask image, the context mask image and the second edge image, the initial model is trained to obtain a trained image restoration model.

[0107] In the embodiment of the present application, steps S401 to S405 can refer to the contents of the aforementioned embodiment and will not be described again here.

[0108] Step S406: Based on the image obtained after preprocessing the captured image during the 3D shooting process, a third edge image is obtained by completing the edge image corresponding to the captured image, and a repair mask image for repairing the captured image is used as a first mask image, and a context mask image for repairing the captured image is used as a second mask image, wherein the first mask image is used to indicate an area to be repaired in the captured image, and the second mask image is used to indicate a known image area when repairing the area to be repaired.

[0109] In an embodiment of the present application, the trained image restoration model can be used for image restoration in 3D photography. It can be understood that in 3D photography, there are usually scenes such as blurring the background. In these scenes, the captured image will be preprocessed to obtain a preprocessed image. The preprocessed image includes known area information for image restoration, the area to be restored, and the completed edge information. Based on this, a third edge image can be obtained after completing the edge image corresponding to the captured image, which is used to restore the restoration mask image of the captured image and used as the first mask image, as well as a context mask image for restoring the captured image and used as the second mask image. Preprocessing may include depth estimation, depth stratification, and mirror movement calculation. In other words, since the known area information, the area to be restored, and the completed edge information used for image restoration are obtained in 3D photography itself, an image for inputting the image restoration model can be generated based on the obtained information. Fig.12 As shown, Fig.12 The second column of images is the second mask image, and the third column of images is the first mask image.

[0110] Step S407: multiplying the second mask image and the captured image to obtain a third image.

[0111] Step S408: Add the third image and the first mask image to obtain a fourth image.

[0112] Step S409: After image connection, the fourth image and the third edge image are input into the image restoration model to obtain a restored region image.

[0113] Similarly, referring to the method of steps S261 to S263 in the above embodiment, the connected image is obtained as the input image, and is input into the image restoration model, so as to obtain the restoration result, that is, the restored region image. Fig.12 As shown, Fig.12 The image in the fourth column is the restored region image, which includes the known image region in the captured image and the restored content of the region to be restored.

[0114] Step S410: based on the repaired regional image, the captured image is repaired to obtain a repaired 3D captured image.

[0115] In the embodiment of the present application, after obtaining the repaired regional image, the image content in the regional image can be embedded into the captured image, so as to complete the repair of the captured image and obtain the final repaired captured image. Since the above-mentioned image repair model is trained based on the completed edge image, it can improve the image repair capability when used for image repair of 3D shooting, and effectively avoid the occurrence of abnormal situations such as broken straight lines and inappropriate content.

[0116] See also Fig.13 , which shows a structural block diagram of a training device 400 for an image restoration model provided in an embodiment of the present application. The training device 400 for the image restoration model uses the above-mentioned computer device, and the training device 400 for the image restoration model includes: a first mask acquisition module 410, a first edge acquisition module 420, a second edge acquisition module 430, a second mask acquisition module 440, and a model training module 450. Among them, the mask acquisition module 410 is used to acquire the mask image corresponding to the sample image, and the mask image is used to indicate the area where the significant object in the sample image is located; the first edge acquisition module 420 is used to extract the edge image of the sample image as the first edge image, and the edge image is used to characterize the texture of the image content of the sample image; the second edge acquisition module 430 is used to complete the edge of the first edge image based on a pre-trained edge completion model, and obtain the completed edge image as the second edge image; the second mask acquisition module 440 is used to acquire the repair mask image and the context mask image corresponding to the sample image based on the mask image, and the repair mask image is used to indicate the area to be repaired in the sample image, and the context mask image is used to indicate the known image area when repairing the area to be repaired; the model training module 450 is used to train the initial model based on the sample image, the repair mask image, the context mask image and the second edge image to obtain the trained image repair model.

[0117] In some embodiments, the model training module 450 can be used to: obtain a background area image based on the repair mask image and the context mask image, wherein the background area image is used to indicate the area to be repaired and the known image area; train the initial model based on the sample image, the repair mask image, the context mask image, the background area image and the second edge image to obtain a trained image repair model.

[0118] As a possible implementation, the model training module 450 trains the initial model based on the sample image, the repair mask image, the context mask image, the background area image and the second edge image to obtain a trained image repair model, which may include: multiplying the background area image with the sample image to obtain a first image; adding the first image with the repair mask image to obtain a second image; performing image connection on the second image and the second edge image to obtain an input image; inputting the input image into the initial model to obtain an output image output by the initial model; based on the repair mask image and the background area image, obtaining a loss value of the output image relative to the sample image as a total loss value; based on the total loss value, iteratively training the initial model until the initial model meets preset conditions to obtain a trained image repair model.

[0119] Optionally, the model training module 450 obtains the loss value of the output image relative to the sample image as a total loss value based on the repair mask image and the background area image, which may include: obtaining the area indicated by the repair mask image from the sample image as a first area image, and obtaining the area indicated by the repair mask image from the output image as a second area image; obtaining the difference between the second area image and the first area image as a first loss value; obtaining the area indicated by the background area image from the sample image as a third area image, and obtaining the area indicated by the background area image from the output image as a fourth area image; obtaining the difference between the fourth area image and the third area image as a second loss value; based on the first loss value and the second loss value, obtaining the loss value of the output image relative to the sample image as a total loss value.

[0120] As a possible implementation, the model training module 450 obtains the background area image based on the repair mask image and the context mask image, which may include: adding the repair mask image and the context mask image to obtain the background area image.

[0121] In some embodiments, the second mask acquisition module 440 can be used to: dilate the mask image to obtain a first processed image; erode the mask image to obtain a second processed image; and based on the mask image, the first processed image and the second processed image, obtain a repaired mask image and a context mask image corresponding to the sample image.

[0122] As a possible implementation, the second mask acquisition module 440 acquires the repaired mask image and the context mask image corresponding to the sample image based on the mask image, the first processed image and the second processed image, which may include: performing a bitwise AND operation on the first processed image and the mask image to obtain the context mask image corresponding to the sample image; performing a bitwise AND operation on the second processed image and the mask image to obtain the repaired mask image corresponding to the sample image.

[0123] As another possible implementation, the second mask acquisition module 440 acquires the repaired mask image and the context mask image corresponding to the sample image based on the mask image, the first processed image and the second processed image, which may include: performing a bitwise AND operation on the first processed image and the mask image to obtain the repaired mask image corresponding to the sample image; performing a bitwise AND operation on the second processed image and the mask image to obtain the context mask image corresponding to the sample image.

[0124] In some embodiments, the second edge acquisition module 430 may be used to: input the first edge image, the sample image, and the mask image into a pre-trained edge completion model to obtain an image output by the edge completion model as the second edge image.

[0125] In some embodiments, the device 400 may further include: a first processing module, a second processing module, a third processing module, a fourth processing module, and an image restoration module. The first processing module is used to obtain a third edge image obtained by completing the edge image corresponding to the captured image based on an image obtained by preprocessing the captured image during the 3D shooting process, and a restoration mask image used to restore the captured image as a first mask image, and a context mask image used to restore the captured image as a second mask image, wherein the first mask image is used to indicate a region to be restored in the captured image, and the second mask image is used to indicate a known image region when restoring the region to be restored; the second processing module is used to multiply the second mask image with the captured image to obtain a third image; the third processing module is used to add the third image with the first mask image to obtain a fourth image; the fourth processing module is used to connect the fourth image and the third edge image, and input them into the image restoration model to obtain a restored region image; the image restoration module is used to restore the captured image based on the restored region image to obtain a restored 3D captured image.

[0126] Those skilled in the art can clearly understand that, for the convenience and brevity of description, the specific working processes of the above-described devices and modules can refer to the corresponding processes in the aforementioned method embodiments, and will not be repeated here.

[0127] In several embodiments provided in the present application, the coupling between modules may be electrical, mechanical or other forms of coupling.

[0128] In addition, each functional module in each embodiment of the present application can be integrated into a processing module, or each module can exist physically separately, or two or more modules can be integrated into one module. The above integrated modules can be implemented in the form of hardware or software functional modules.

[0129] In summary, the solution provided by the present application is to obtain a mask image corresponding to a sample image, the mask image is used to indicate the area where a significant object is located in the sample image, extract the edge image of the sample image as the first edge image, the edge image is used to characterize the texture of the image content of the sample image, and based on a pre-trained edge completion model, the first edge image is edge-completed to obtain the completed edge image as the second edge image, and the initial model is trained based on the sample image, the mask image and the second edge image to obtain a trained image restoration model. Since the extracted edge image is edge-completed during model training, and the area to be restored and the known image area are refined based on the mask image, the trained model can better restore the image content and improve the image restoration effect.

[0130] Please refer to Fig.14 , which shows a block diagram of a computer device provided in an embodiment of the present application. The computer device 100 may be a computer device capable of running applications, such as a server, a laptop, a PC, a smart phone, a tablet computer, etc. The computer device 100 in the present application may include one or more of the following components: a processor 110, a memory 120, and one or more applications, wherein the one or more applications may be stored in the memory 120 and configured to be executed by one or more processors 110, and the one or more programs are configured to execute the method described in the aforementioned method embodiment.

[0131] The processor 110 may include one or more processing cores. The processor 110 uses various interfaces and lines to connect various parts of the entire computer device 100, and executes various functions and processes data of the computer device 100 by running or executing instructions, programs, code sets or instruction sets stored in the memory 120, and calling data stored in the memory 120. Optionally, the processor 110 can be implemented in at least one hardware form of digital signal processing (DSP), field-programmable gate array (FPGA), and programmable logic array (PLA). The processor 110 can integrate one or a combination of a central processing unit (CPU), a graphics processing unit (GPU), and a modem. Among them, the CPU mainly processes the operating system, user interface, and application programs; the GPU is responsible for rendering and drawing display content; and the modem is used to process wireless communications. It can be understood that the above-mentioned modem may not be integrated into the processor 110, but may be implemented separately through a communication chip.

[0132] The memory 120 may include a random access memory (RAM) or a read-only memory (ROM). The memory 120 may be used to store instructions, programs, codes, code sets or instruction sets. The memory 120 may include a program storage area and a data storage area, wherein the program storage area may store instructions for implementing an operating system, instructions for implementing at least one function (such as a touch function, a sound playback function, an image playback function, etc.), instructions for implementing the following various method embodiments, etc. The data storage area may also store data (such as a phone book, audio and video data, chat record data) created by the computer device 100 during use.

[0133] Please refer to Fig.15 , which shows a structural block diagram of a computer-readable storage medium provided in an embodiment of the present application. The computer-readable medium 800 stores program codes, which can be called by a processor to execute the method described in the above method embodiment.

[0134] The computer readable storage medium 800 may be an electronic memory such as a flash memory, an EEPROM (electrically erasable programmable read-only memory), an EPROM, a hard disk, or a ROM. Optionally, the computer readable storage medium 800 includes a non-transitory computer-readable storage medium. The computer readable storage medium 800 has storage space for program code 810 that performs any method step in the above method. These program codes can be read from or written to one or more computer program products. The program code 810 can be compressed, for example, in an appropriate form.

[0135] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present application, rather than to limit it. Although the present application has been described in detail with reference to the aforementioned embodiments, those skilled in the art should understand that they can still modify the technical solutions described in the aforementioned embodiments, or make equivalent replacements for some of the technical features therein. However, these modifications or replacements do not cause the essence of the corresponding technical solutions to deviate from the spirit and scope of the technical solutions of the embodiments of the present application.

Claims

1. An image restoration method, characterized in that: The method comprises: Acquire a mask image corresponding to the sample image, where the mask image is used to indicate an area where a significant object in the sample image is located; Extracting an edge image of the sample image as a first edge image, wherein the edge image is used to characterize a texture of an image content of the sample image; Based on a pre-trained edge completion model, edge completion is performed on the first edge image to obtain a completed edge image as a second edge image; Based on the mask image, a repair mask image and a context mask image corresponding to the sample image are obtained, wherein the repair mask image is used to indicate an area to be repaired in the sample image, and the context mask image is used to indicate a known image area when repairing the area to be repaired, the repair mask image is obtained by performing a bitwise AND operation on the mask image and the mask image after a first process, and the context mask image is obtained by performing a bitwise AND operation on the mask image and the mask image after a second process, the first process and the second process are dilation processes or erosion processes, and the first process and the second process have different processing types; Based on the sample image, the restoration mask image, the context mask image and the second edge image, an initial model is trained to obtain a trained image restoration model.

2. The method according to claim 1, characterized in that The initial model is trained based on the sample image, the restoration mask image, the context mask image and the second edge image to obtain a trained image restoration model, and the trained image restoration model is obtained, including: Based on the repair mask image and the context mask image, a background area image is obtained, wherein the background area image is used to indicate the area to be repaired and the known image area; Based on the sample image, the restoration mask image, the context mask image, the background area image and the second edge image, an initial model is trained to obtain a trained image restoration model.

3. The method according to claim 2, characterized in that The step of training the initial model based on the sample image, the restoration mask image, the context mask image, the background area image, and the second edge image to obtain a trained image restoration model includes: Multiplying the background area image and the sample image to obtain a first image; Adding the first image to the repair mask image to obtain a second image; Connecting the second image and the second edge image to obtain an input image; Inputting the input image into the initial model to obtain an output image output by the initial model; Based on the repair mask image and the background area image, obtaining a loss value of the output image relative to the sample image as a total loss value; Based on the total loss value, the initial model is iteratively trained until the initial model meets a preset condition, thereby obtaining a trained image restoration model.

4. The method according to claim 3, characterized in that The acquiring, based on the repair mask image and the background area image, a loss value of the output image relative to the sample image as a total loss value comprises: Acquire the region indicated by the repair mask image from the sample image as a first region image, and acquire the region indicated by the repair mask image from the output image as a second region image; obtaining a difference between the second region image and the first region image as a first loss value; Acquire the area indicated by the background area image from the sample image as a third area image, and acquire the area indicated by the background area image from the output image as a fourth area image; obtaining a difference between the fourth region image and the third region image as a second loss value; Based on the first loss value and the second loss value, a loss value of the output image relative to the sample image is obtained as a total loss value.

5. The method according to claim 2, characterized in that: The step of obtaining a background area image based on the repair mask image and the context mask image includes: The repair mask image is added to the context mask image to obtain a background area image.

6. The method according to claim 1, characterized in that The step of acquiring a repaired mask image and a context mask image corresponding to the sample image based on the mask image includes: Performing dilation processing on the mask image to obtain a first processed image; Performing corrosion processing on the mask image to obtain a second processed image; Based on the mask image, the first processed image, and the second processed image, a repair mask image and a context mask image corresponding to the sample image are acquired.

7. The method according to claim 6, characterized in that The acquiring, based on the mask image, the first processed image, and the second processed image, a repaired mask image and a context mask image corresponding to the sample image includes: Performing a bitwise AND operation on the first processed image and the mask image to obtain a context mask image corresponding to the sample image; Performing a bitwise AND operation on the second processed image and the mask image to obtain a repaired mask image corresponding to the sample image; or The acquiring, based on the mask image, the first processed image, and the second processed image, a repaired mask image and a context mask image corresponding to the sample image includes: Performing a bitwise AND operation on the first processed image and the mask image to obtain a repaired mask image corresponding to the sample image; A bitwise AND operation is performed on the second processed image and the mask image to obtain a context mask image corresponding to the sample image.

8. The method according to claim 1, characterized in that The step of performing edge completion on the first edge image based on the pre-trained edge completion model to obtain the completed edge image as the second edge image includes: The first edge image, the sample image, and the mask image are input into a pre-trained edge completion model, and an image output by the edge completion model is obtained as the second edge image.

9. The method according to any one of claims 1 to 8, characterized in that: After the initial model is trained based on the sample image, the mask image and the second edge image to obtain a trained image restoration model, the method further includes: Based on an image obtained after preprocessing a captured image during 3D shooting, a third edge image obtained by completing an edge image corresponding to the captured image is acquired, a repair mask image used to repair the captured image is used as a first mask image, and a context mask image used to repair the captured image is used as a second mask image, wherein the first mask image is used to indicate a to-be-repaired area in the captured image, and the second mask image is used to indicate a known image area when repairing the to-be-repaired area; multiplying the second mask image by the captured image to obtain a third image; adding the third image to the first mask image to obtain a fourth image; After image connection, the fourth image and the third edge image are input into the image restoration model to obtain a restored region image; The captured image is restored based on the restored regional image to obtain a restored 3D captured image.

10. A training device for an image restoration model, characterized in that: The device comprises: a first mask acquisition module, a first edge acquisition module, a second edge acquisition module, a second mask acquisition module and a model training module, wherein: The mask acquisition module is used to acquire a mask image corresponding to a sample image, wherein the mask image is used to indicate an area where a significant object in the sample image is located; The first edge acquisition module is used to extract an edge image of the sample image as a first edge image, wherein the edge image is used to characterize the texture of the image content of the sample image; The second edge acquisition module is used to perform edge completion on the first edge image based on a pre-trained edge completion model to obtain a completed edge image as a second edge image; The second mask acquisition module is used for acquiring a repair mask image and a context mask image corresponding to the sample image based on the mask image, the repair mask image is used for indicating an area to be repaired in the sample image, and the context mask image is used for indicating a known image area when repairing the area to be repaired, the repair mask image is obtained by performing a bitwise AND operation on the mask image and the mask image after the first processing, and the context mask image is obtained by performing a bitwise AND operation on the mask image and the mask image after the second processing, the first processing and the second processing are dilation processing or erosion processing, and the processing types of the first processing and the second processing are different; The model training module is used to train the initial model based on the sample image, the restoration mask image, the context mask image and the second edge image to obtain a trained image restoration model.

11. A computer device, characterized in that: include: one or more processors; Memory; One or more application programs, wherein the one or more application programs are stored in the memory and configured to be executed by the one or more processors, and the one or more programs are configured to execute the method according to any one of claims 1-9.

12. A computer-readable storage medium, characterized in that: The computer-readable storage medium stores program codes, which can be called by a processor to execute the method according to any one of claims 1 to 9.

Images