Image restoration method, system and electronic device

Abstract

The present invention discloses an image restoration method, system and electronic device, and relates to the field of image restoration technology. The present invention divides the traditional image restoration method into a detection stage of an area to be restored and an image restoration stage. First, a target detection model is used to detect the area to be restored, and then only the area to be restored is input into the image restoration model for image restoration, which can avoid the waste of resources caused by restoring the entire image and improve the restoration efficiency. In addition, the present invention trains an improved yolov5 network to obtain a target detection model, which can solve the problem of poor detection accuracy or false detection and missed detection of the original yolov5 network under complex backgrounds and multi-noise conditions, and trains an improved GAN network to obtain an image restoration model, which can solve the problem that the traditional GAN ​​network has a poor restoration effect on large-area damaged images, thereby improving the image restoration effect.

Description

Technical Field

[0001] The present invention relates to the technical field of image restoration, and in particular to an image restoration method, system and electronic equipment. Background Art

[0002] With the rapid development of the information age and the emergence of various image sensors, a large amount of data is transmitted in the form of images on the Internet. Due to partial occlusion of the sensor or pixel loss during the transmission process, the image information received by the image receiving device is lost and the image is incomplete. Therefore, image restoration technology is used to repair the missing parts of the image and enhance the image quality.

[0003] For traditional image restoration algorithms, most methods can only obtain useful information from the input image and repair the damaged or occluded image through diffusion or block filling methods. However, if the image is damaged over a large area, the restoration effect of the traditional image restoration algorithm will be very poor, and it will almost lose the restoration ability. Moreover, image data is structural data, and traditional image restoration algorithms lack the ability to perceive contextual semantic information, so the restoration effect is often poor.

[0004] With the advancement of deep learning, image restoration based on deep learning has gradually grown up, relying on the "hallucination" of trained neural networks to fill the large holes in the image; especially now the rapid development of convolutional neural networks and the design of generative adversarial networks have enabled the field of image processing to develop rapidly, and image restoration technology based on deep learning has also been widely and deeply studied.

[0005] Nowadays, Generative Adversarial Networks (GAN) has been widely studied and is a hot topic in image restoration technology. The traditional image restoration method is to directly input the entire image into the restoration network for restoration. This restoration method will cause unnecessary restoration of the undamaged image area, which will cause computational pressure on the restoration network and cause pixel loss, and the restoration efficiency is not high. In addition, due to the imperfect network structure of the traditional GAN ​​network, the restoration results of large-area defective images are relatively fuzzy and the discriminator has poor identification ability. Summary of the invention

[0006] In order to solve the above problems existing in the prior art, the present invention provides an image restoration method, system and electronic device.

[0007] To achieve the above object, the present invention provides the following solutions:

[0008] An image restoration method, comprising:

[0009] Obtain object detection datasets and image restoration datasets;

[0010] The target detection data set is used to train the improved yolov5 network to obtain a target detection model;

[0011] The image restoration data set is used to train the improved GAN network to obtain an image restoration model;

[0012] Inputting the image to be repaired into the target detection model to detect the area to be repaired;

[0013] Inputting the area to be repaired into the image repair model to obtain an image repair result;

[0014] The image restoration result and the image to be restored are fused to obtain a restored image.

[0015] Optionally, a box attention mechanism is added to the yolov5 network to obtain the improved yolov5 network.

[0016] Optionally, in the yolov5 network, when the deep features are fused with the shallow features of the next layer, the box attention mechanism is added to obtain the improved yolov5 network.

[0017] Optionally, the improved GAN network is obtained after introducing a collaborative modulation mechanism and a dual contrast loss function into the GAN network.

[0018] Optionally, the process of introducing the cooperative modulation mechanism in the GAN network includes: modulating the unconditional coordination generator into a cooperative modulation generator;

[0019] In the process of modulating to a co-modulation generator, the latent factors are transformed through eight fully connected networks and affine transformed to generate a style vector.

[0020] Optionally, the dual contrast loss function is

[0021] In the formula, G is the generator, D is the discriminator, is the contrast loss function in the first case, is the contrast loss function in the second case, Minimize the probability that the discriminator can distinguish fake samples generated by the generator as real samples. Maximize the probability that the discriminator identifies the true sample as the true sample.

[0022] According to the specific embodiments provided by the present invention, the present invention discloses the following technical effects:

[0023] The image restoration method provided by the present invention divides the traditional image restoration method into a detection stage of the area to be restored and an image restoration stage. First, the area to be restored is detected by using a target detection model, and then only the area to be restored is input into the image restoration model for image restoration, which can avoid the waste of resources caused by restoring the entire image and improve the restoration efficiency. In addition, the present invention trains the improved yolov5 network to obtain a target detection model, which can solve the problem of poor detection accuracy or false detection and missed detection of the original yolov5 network under complex backgrounds and multi-noise conditions, and trains the improved GAN network to obtain an image restoration model, which can solve the problem that the traditional GAN ​​network has a poor restoration effect on large-area damaged images, thereby improving the image restoration effect.

[0024] Furthermore, the present invention also provides the following implementation structure:

[0025] An image restoration system is applied to the image restoration method provided above; the system comprises:

[0026] The data set acquisition module is used to acquire the target detection data set and the image restoration data set;

[0027] A target detection model training module, used for training an improved yolov5 network using the target detection data set to obtain a target detection model; the improved yolov5 network is a network obtained by adding a box attention mechanism to the yolov5 network;

[0028] An image restoration model training module, used to train an improved GAN network using the image restoration data set to obtain an image restoration model; the improved GAN network is a network obtained by introducing a collaborative modulation mechanism and a dual contrast loss function into the GAN network;

[0029] A to-be-repaired area detection module is used to input the to-be-repaired image into the target detection model to detect the to-be-repaired area;

[0030] A restoration result determination module, used for inputting the area to be restored into the image restoration model to obtain an image restoration result;

[0031] The image restoration module is used to fuse the image restoration result and the image to be restored to obtain a restored image.

[0032] An electronic device, comprising:

[0033] Memory for storing computer programs;

[0034] A processor is connected to the memory and is used to retrieve and execute the computer program to implement the image restoration method provided above.

[0035] A computer-readable storage medium is used to store a computer program; the computer program is used to implement the image restoration method provided above.

[0036] Since the technical effect achieved by the implementation structure provided by the present invention is the same as the technical effect achieved by the image restoration method provided by the present invention, it will not be described in detail here. BRIEF DESCRIPTION OF THE DRAWINGS

[0037] In order to more clearly illustrate the embodiments of the present invention or the technical solutions in the prior art, 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.

[0038] Figure 1 A flowchart of the image restoration method provided by the present invention;

[0039] Figure 2 A schematic diagram of the yolov5 network structure of the fusion box attention mechanism provided by the present invention;

[0040] Figure 3 A schematic diagram of the structure of the box attention mechanism provided by the present invention;

[0041] Figure 4 The modulation provided by the present invention is a schematic diagram of the structure of a collaborative modulation generator; wherein, Figure 4 (a) is a schematic diagram of the structure of the unconditional modulation generator. Figure 4 (b) is a schematic diagram of the structure of the condition generator. Figure 4 (c) is a schematic diagram of the structure of the conditional modulation generator. Figure 4 (d) is a schematic diagram of the structure of the cooperative modulation generator;

[0042] Figure 5 A schematic diagram of an unconditional coordination generator provided by the present invention. DETAILED DESCRIPTION

[0043] The following will be combined with the drawings in the embodiments of the present invention to clearly and completely describe the technical solutions in the embodiments of the present invention. Obviously, the described embodiments are only part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without creative work are within the scope of protection of the present invention.

[0044] The purpose of the present invention is to provide an image restoration method, system and electronic equipment, which can improve the image restoration effect and image restoration efficiency.

[0045] In order to make the above-mentioned objects, features and advantages of the present invention more obvious and easy to understand, the present invention is further described in detail below with reference to the accompanying drawings and specific embodiments.

[0046] like Figure 1 As shown, the image restoration method provided by the present invention comprises:

[0047] Step 100: Obtain target detection dataset and image restoration dataset. Before this step, different target detection datasets and image restoration datasets need to be produced according to the different objects to be restored. Among them, the target detection dataset of the area to be restored is collected according to its own needs, and the data is labeled.

[0048] Step 101: Use the target detection data set to train the improved yolov5 network to obtain the target detection model. This step mainly uses the BA-YOLO target detection algorithm to train the target detection model for the area to be repaired. Specifically, the BA-YOLO algorithm that integrates the box attention mechanism into the feature fusion process of the yolov5 network is used to train the target detection data set. Among them, the yolov5 network structure with the fusion box attention mechanism is as follows Figure 2 shown. Figure 2 The Chinese and English translations of the module names are shown in Table 1.

[0049] Table 1 Module name Chinese-English comparison table

[0050]

[0051]

[0052] In practical applications, BoxAttention belongs to a multi-head attention mechanism. BoxAttention aggregates multiple heads with the same features and combines them with a learnable projection matrix W. O Attention calculation is performed based on multi-head self-attention. Given the query (transliterated as intent) vector q in the i-th attention head, the interest box is b i , BoxAttention uses bilinear interpolation from the box of interest b i Extract the feature map v of size m×m i In the calculation of the i-th head attention, the feature map v i is a set of values ​​corresponding to the query vector q, and then calculates the query vector q and the m×m learnable Key vector K iThe dot product between them is passed through the softmax function to obtain the attention weight map, and finally multiplied by the value V of the corresponding position to obtain the feature map v with attention weight size of m×m i The weighted average h i for:

[0053]

[0054] In the formula, Q, K i ,V represent the query volume, key value and value respectively. BoxAttention(Q,K i ,V i ) is the attention head mechanism, V i is the i-th value.

[0055] According to the weighted average h obtained above i , and then when the features are fused, Figure 3 As shown in the figure, when the deep features are fused with the features of the next shallow layer, the attention mechanism is added to allocate weights, which is equivalent to finding a region of interest ROI in the small-scale large target information, and then using the shallow information to supplement the details and suppress irrelevant information. Therefore, adding the box attention mechanism to the yolov5 network can enhance the ability to extract key features, thereby enhancing the feature fusion effect of the network, reducing the impact of noise and complex background, and improving the detection accuracy of the target (i.e., the area to be repaired), and can more accurately find the target in the area to be repaired in the image.

[0056] Step 102: Use the image restoration dataset to train the improved GAN network to obtain an image restoration model. This step mainly uses a GAN model that introduces a collaborative modulation mechanism and a dual contrast loss function to train the image restoration model.

[0057] In the actual application process, in order to enhance the ability to repair large-area defective images, a collaborative modulation mechanism is introduced into the GAN network to combine the advantages of the conditional GAN ​​network and the unconditional GAN ​​network. Figure 4 (a) shows that the process of modulation of the latent factor Z is as follows Figure 4 (a) to Figure 4 (d) and Figure 5As shown in the figure, the latent factor Z is transformed into an intermediate vector W by eight fully connected networks M. At this time, the intermediate vector W can make the attribute distribution more uniform while maintaining consistency with the training set, and maintain a better linear relationship with the generated image, which is conducive to the attribute control of the generated image. Therefore, the intermediate vector W is more suitable as the input of the generator. After that, the intermediate vector W is affine transformed A and then demodulated to obtain the style that can control each layer of the synthetic network. In addition, additional noise can help the synthetic network improve the performance of random transformation. The style vector s generated by the latent factor Z after the mapping of the eight fully connected networks M and the affine transformation A can be expressed as:

[0058] s=A(M(Z)) (2)

[0059] Where M(Z) represents the mapping of the latent factor Z by eight fully connected networks, and A(·) represents the affine transformation.

[0060] The co-modulation generator can coordinate the differences between the conditional GAN ​​network and the unconditional GAN ​​network, naturally learning and utilizing random styles, such as Figure 4 As shown in (d). At this time, the style vector s can be expressed as:

[0061] s=A(E(Y),M(Z)) (3)

[0062] In the formula, E(Y) represents the condition generator and Y represents the condition input.

[0063] In order to improve the discriminative ability of the discriminator in the GAN network, a dual contrast loss function is used based on the introduction of the collaborative modulation mechanism. In the first case, a real picture is added to all the generated images, and the discriminator is required to select this real image. The contrast loss function is:

[0064]

[0065] In the formula, is the expected function, p(x) represents the distribution of the real sample, N(0,I d ) represents the distribution of fake samples, I d represents the range of false samples, D(x) represents the probability that the generative discriminator classifies sample x as real data, and D(G(z)) represents the probability of distinguishing a false sample from a real sample.

[0066] In the second case, when a generated image is placed among a batch of real images and the discriminator is required to identify the generated image, the contrast loss function is

[0067]

[0068] The training process in these two scenarios is expressed as a unified dual contrast loss function:

[0069]

[0070] In the formula, G is the generator, D is the discriminator, is the contrast loss function in the first case, is the contrast loss function in the second case, Minimize the probability that the discriminator can distinguish fake samples generated by the generator as real samples. Maximize the probability that the discriminator identifies the true sample as the true sample.

[0071] On the basis of collaborative modulation, the dual contrast loss function is introduced, which enables the discriminator to select the real image from multiple generated images and one real image, and then identify the generated image from multiple real images and one generated image, so as to enhance the discriminator's ability to distinguish and enable it to compete with the generator. Combining these two mechanisms can.

[0072] Through this contrastive learning, the loss function of the GAN network can strengthen the discriminator's ability to distinguish between real images and generated images, and can have sufficient generalization ability to stimulate the confrontation with the generator (that is, the generative network in the GAN network), transmit more meaningful information to the generator, promote the evolution of the generator, make the samples generated by the generator semantically coherent, and enhance the model's image restoration ability. In addition, it can also enhance the restoration ability and restoration effect of large-area defective images.

[0073] Step 103: Input the image to be repaired into the target detection model to detect the area to be repaired.

[0074] Step 104: input the area to be repaired into the image repair model to obtain an image repair result.

[0075] Step 105: Fusing the image restoration result and the image to be restored to obtain a restored image, that is, fusing the output image restoration result with the background to form an image with a size equivalent to the image to be restored, which is the restored image.

[0076] Based on the above description, compared with the prior art, the present invention has the following characteristics:

[0077] i) The traditional image restoration method is decomposed into two stages: detection of the area to be restored and restoration of the defective image, which improves the efficiency of image restoration.

[0078] ii) The box attention mechanism is introduced into the YOLOv5 network, and the BA-YOLO algorithm is proposed, which enhances the attention of the YOLOv5 network to key features, ignores the influence of background and noise, and improves the object detection accuracy in the area to be repaired.

[0079] iii) The collaborative modulation mechanism and dual contrast loss function are introduced simultaneously in the traditional GAN ​​network, which can enhance the model's repair effect on large-area defective images and improve the ability of the discriminator.

[0080] Furthermore, the present invention also provides the following implementation structure:

[0081] An image restoration system is applied to the image restoration method provided above. The system comprises:

[0082] The dataset acquisition module is used to obtain target detection datasets and image restoration datasets.

[0083] The target detection model training module is used to train the improved yolov5 network using the target detection dataset to obtain the target detection model. The improved yolov5 network is a network obtained by adding the box attention mechanism to the yolov5 network.

[0084] The image restoration model training module is used to train the improved GAN network using the image restoration dataset to obtain the image restoration model. The improved GAN network is a network obtained by introducing the collaborative modulation mechanism and the dual contrast loss function into the GAN network.

[0085] The to-be-repaired area detection module is used to input the to-be-repaired image into the target detection model to detect the to-be-repaired area.

[0086] The repair result determination module is used to input the area to be repaired into the image repair model to obtain the image repair result.

[0087] The image restoration module is used to fuse the image restoration result and the image to be restored to obtain the restored image.

[0088] An electronic device, comprising:

[0089] Memory, used to store computer programs.

[0090] The processor is connected to the memory and is used to retrieve and execute the computer program to implement the image restoration method provided above.

[0091] In addition, when the computer program in the above-mentioned memory is implemented in the form of a software functional unit and sold or used as an independent product, it can be stored in a computer-readable storage medium. Based on this understanding, the technical solution of the present invention, or the part that contributes to the prior art or the part of the technical solution, can be embodied in the form of a software product, which is stored in a storage medium and includes several instructions for a computer device (which can be a personal computer, a server or a network device, etc.) to perform all or part of the steps of the methods of each embodiment of the present invention. The aforementioned storage medium includes: various media that can store program codes, such as a USB flash drive, a mobile hard disk, a read-only memory, a random access memory, a magnetic disk or an optical disk.

[0092] The present invention divides the traditional image restoration method into a to-be-restored area detection stage and an image restoration stage. First, the to-be-restored area is detected, and then only the to-be-restored area is input into the image restoration network, thereby avoiding the waste of resources caused by repairing the entire image and improving the restoration efficiency.

[0093] In addition, in the detection stage of the area to be repaired, the present invention adds the box attention mechanism BoxAttention method to the yolov5 network and proposes the BA-YOLO algorithm to solve the problem of poor detection accuracy or false detection and missed detection of the original yolov5 network under complex backgrounds and multi-noise conditions. The box attention mechanism can make the yolov5 network pay more attention to key features when performing feature fusion and ignore the influence of background and noise. The BA-YOLO algorithm proposed in the present invention has higher detection accuracy for the target area to be repaired than the detection accuracy of the yolov5 network.

[0094] In the image restoration stage, the present invention proposes an algorithm that introduces a collaborative modulation mechanism and a dual contrast loss function in a GAN network to optimize the image restoration performance. Aiming at the problem that the traditional GAN ​​network has a poor restoration effect on large-area damaged images, the present invention introduces a collaborative modulation mechanism in the GAN network to combine the advantages of the conditional GAN ​​network in generating controllable image styles and the advantages of the non-conditional GAN ​​network in generating images randomly, thereby improving the restoration performance of the restoration network for large-area damaged images. At the same time, aiming at the shortcomings that the loss function of the traditional GAN ​​is insufficiently generalized for the feature representation of the discriminator and cannot stimulate the adversarial evolution of the generator, the present invention uses a dual contrast loss function in the GAN network on the basis of introducing a collaborative modulation mechanism, thereby enhancing the identification ability of the discriminator and enabling it to compete with the generator. The proposed image restoration algorithm that introduces a collaborative modulation mechanism and a dual contrast loss function has better effects and better performance in large-area image restoration work.

[0095] In this specification, each embodiment is described in a progressive manner, and each embodiment focuses on the differences from other embodiments. The same or similar parts between the embodiments can be referred to each other. For the system disclosed in the embodiment, since it corresponds to the method disclosed in the embodiment, the description is relatively simple, and the relevant parts can be referred to the method part.

[0096] The principles and implementation methods of the present invention are described in this article using specific examples. The description of the above embodiments is only used to help understand the method and core idea of ​​the present invention. At the same time, for those skilled in the art, according to the idea of ​​the present invention, there will be changes in the specific implementation methods and application scope. In summary, the content of this specification should not be understood as limiting the present invention.

Claims

1. An image restoration method, characterized in that: include: Obtain object detection datasets and image restoration datasets; The target detection data set is used to train the improved yolov5 network to obtain the target detection model; wherein, in the yolov5 network, a box attention mechanism is added when the deep-level features are fused with the shallow-level features of the next layer to obtain the improved yolov5 network; The image restoration data set is used to train the improved GAN network to obtain an image restoration model; wherein the improved GAN network is obtained after introducing a collaborative modulation mechanism and a dual contrast loss function into the GAN network; Inputting the image to be repaired into the target detection model to detect the area to be repaired; Inputting the area to be repaired into the image repair model to obtain an image repair result; Fusing the image restoration result and the image to be restored to obtain a restored image; The process of introducing the collaborative modulation mechanism into the GAN network includes: modulating the unconditional coordinated generator into a collaborative modulation generator; in the process of modulating into the collaborative modulation generator, the latent factor generates a style vector after being transformed by eight fully connected networks and affine transformation; The dual contrast loss function is In the formula, G is the generator, D is the discriminator, is the contrast loss function in the first case, is the contrast loss function in the second case, Minimize the probability that the discriminator can distinguish fake samples generated by the generator as real samples. The discriminator is designed to maximize the probability of identifying a true sample as a true sample. The first case refers to adding a real picture to all generated images, and the discriminator is required to select this real image. The second case refers to putting a generated picture in a batch of real pictures, and the discriminator is required to identify this generated picture.

2. An image restoration system, characterized in that: Applied to the image restoration method as claimed in claim 1; the system comprises: The data set acquisition module is used to acquire the target detection data set and the image restoration data set; A target detection model training module, used for training an improved yolov5 network using the target detection data set to obtain a target detection model; the improved yolov5 network is a network obtained by adding a box attention mechanism to the yolov5 network; An image restoration model training module, used to train an improved GAN network using the image restoration data set to obtain an image restoration model; the improved GAN network is a network obtained by introducing a collaborative modulation mechanism and a dual contrast loss function into the GAN network; A to-be-repaired area detection module is used to input the to-be-repaired image into the target detection model to detect the to-be-repaired area; A restoration result determination module, used for inputting the area to be restored into the image restoration model to obtain an image restoration result; The image restoration result and the image to be restored are fused to obtain a restored image.

3. An electronic device, characterized in that: include: Memory for storing computer programs; A processor is connected to the memory and is used to retrieve and execute the computer program to implement the image restoration method as claimed in claim 1.

4. A computer-readable storage medium, characterized in that: The computer-readable storage medium is used to store a computer program; the computer program is used to implement the image restoration method as claimed in claim 1.

Images