Image restoration data set distillation method based on uniform image selection and neural representation

Abstract

The invention relates to the field of machine learning and image processing, provides an image restoration data set distillation method based on uniform image selection and neural representation, and solves the problems of large calculation amount, poor distillation performance and the like of an existing matching-based method. The method comprises the steps of training a restoration model in an original data set; based on the restoration model, calculating peak signal-to-noise ratios of restoration results of the low-quality samples in the training set and sorting the peak signal-to-noise ratios; removing partial images with the lowest peak signal-to-noise ratio and the highest peak signal-to-noise ratio, and uniformly selecting the images according to the compression ratio; respectively training the corresponding neural representation of each selected image; the mask nerve represents a part of pixels with the minimum reconstruction loss; reconstructing a recovery data set based on the mask coordinates, the distillation data and the neural representation; and training a deep learning model on the reconstructed data and testing the restoration performance. The method has the advantages that the image restoration model is efficiently trained, the storage space of an image restoration data set is greatly reduced, and iterative updating of a restoration algorithm is accelerated.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of machine learning and image processing, and particularly relates to an image restoration dataset distillation method based on uniform image selection and neural representation. BACKGROUND

[0002] In the past decade, with the rapid development of deep learning, researchers have proposed more and more large-scale datasets for different downstream tasks, and have trained many excellent models based on these datasets. However, these large-scale datasets have led to an increase in storage space and training time, high training costs, user information privacy leakage and other problems, so that the dataset distillation direction has gradually become a research hotspot. Given an original dataset, dataset distillation aims to generate a smaller dataset, so that the training effect of the same model on the small dataset can be comparable to that on the original dataset. Dataset distillation combines multiple fields such as machine learning and image processing, and can accelerate the algorithm training and iteration of deep learning and effectively reduce the storage space of the dataset.

[0003] Existing dataset distillation methods mainly include matching-based methods, generative model-based methods and non-optimization methods. The matching-based method includes two-layer optimization performance matching and gradient matching method, in which the outer loop optimizes the image pixels in the dataset, and the inner loop optimizes the model. In addition, it also includes single-layer optimization distribution matching, batch regularization matching method, etc. This method needs to optimize image pixels, which is computationally intensive. The generative model-based method fine-tunes the pre-trained generative model by introducing a new loss function, so that it generates the most representative samples, but this method depends on the pre-trained generative model. The non-optimization method calculates the loss value of each sample in the training set through the pre-trained model of the downstream task, and selects the image block with the lowest loss as the distillation sample. Subsequently, these samples are input into the pre-trained model to obtain the corresponding soft label, thereby generating the final distillation data. However, the distribution of the distillation data obtained by this method is different from that of the original dataset. Most existing dataset distillation researches focus on image classification tasks. The dataset distillation method for image restoration tasks is relatively scarce, and the current main method is the image super-resolution distillation method based on difficult sample selection. However, this method performs poorly in tasks such as severe weather image restoration, image deblurring and image denoising. In addition, since the image restoration task focuses on texture details, and the image classification focuses on semantic information, if the dataset distillation method for image classification is directly migrated to the restoration task, the texture details of the distilled image will be lost, and the restoration performance will be greatly reduced. Therefore, it is urgent to propose a new framework for image restoration tasks. SUMMARY

[0004] In view of the problems of large amount of calculation and poor performance of the existing matching-based distillation method, the application provides an image restoration dataset distillation method based on uniform image selection and neural representation. The method first obtains the restoration difficulty of each sample in the training set based on a pre-trained restoration model, removes extreme samples, and then selects images uniformly according to the compression ratio. In order to further improve the performance, the neural representation corresponding to each image is trained, and the pixels with the lowest mask reconstruction loss based on the neural representation are obtained to obtain the distillation dataset. In the deployment stage, the neural representation and the mask coordinates are used to reconstruct the restoration dataset. The performance of the distillation algorithm is evaluated by training the restoration model on the reconstructed dataset. The whole method can be applied to images of different resolutions, and the effectiveness of the proposed distillation method is verified in six image restoration tasks of raindrop removal, rain and fog joint removal, defogging, snow removal, deblurring and denoising.

[0005] The object of the application can be achieved by the following technical solutions:

[0006] An image restoration dataset distillation method based on uniform image selection and neural representation, characterized in that the method comprises the following steps:

[0007] Step 1: training a deep learning model on an image restoration dataset, and giving a final dataset compression ratio r, wherein 0

[0008] Step 2: inputting the low-quality samples in each training set into the model, calculating the peak signal-to-noise ratio between the restoration result and the real sample, and sorting from small to large.

[0009] Step 3: removing the images with the lowest and highest peak signal-to-noise ratio in the proportion of t, and selecting images uniformly according to the compression ratio r1, wherein 0

[0010] Step 4: training the neural representation corresponding to each selected image, that is, the input of the neural network is the coordinates, and the output is the RGB pixel value of the corresponding coordinates.

[0011] Step 5: reconstructing the pixels with the lowest image loss based on the neural representation for each selected image mask, and recording the mask coordinates as the final distillation dataset.

[0012] Step 6: reconstructing the restoration dataset based on the mask coordinates and the neural representation.

[0013] Step 7: training the deep learning model on the reconstructed data, calculating the peak signal-to-noise ratio and structural similarity in the test set, and evaluating the performance of the distillation algorithm.

[0014] ​Further, the image restoration dataset distillation method based on uniform image selection and neural representation has the characteristics that the deep learning model trained on the restoration dataset in step 1 includes an encoder and a decoder, which are used for extracting low-quality image features and generating restored images, respectively.

[0015] Further, the image restoration dataset distillation method based on uniform image selection and neural representation has the characteristics that the specific expression of the peak signal-to-noise ratio in step 2 is as follows:

[0016]

[0017] wherein w and h respectively represent the width and height of the training image, x and y respectively represent the paired low-quality image and real image, f θ represents the restoration model trained in step 1, and the training parameter is θ. The peak signal-to-noise ratio as the image restoration difficulty, the higher the value, the easier the restoration, and the lower the value, the more difficult the restoration.

[0018] Further, the image restoration dataset distillation method based on uniform image selection and neural representation has the characteristics that in step 3, considering that the subsequent neural representation will mask part of the pixels, therefore, the proportion r1 of the selected image will be larger than the pre-specified compression ratio r. If the low-quality-real image sequence sorted from small to large peak signal-to-noise ratio is {x1, y1}, {x2, y2}, … {x n ,y n}, then the selected image is:

[0019]

[0020] wherein n represents the number of images of the original dataset, r1 is the proportion of the selected image in the first stage, {x j ,y j} represents the low-quality image and real image pair ranked j from small to large peak signal-to-noise ratio, t is the proportion of images excluding the highest and lowest peak signal-to-noise ratio, and satisfies 0<t<1, represents the floor operator.

[0021] Further, the image restoration dataset distillation method based on uniform image selection and neural representation has the characteristics that in step 4, for an image with a width W and a height H, the neural representation loss function of the image is:

[0022]

[0023] wherein, The output of the neural representation with the given coordinate (w, h) and the parameter θ1, x(w, h) represents the RGB pixel value of the image coordinate (w, h). The low-quality image and the real image in step 3 are both involved in training. Since only the pixels of the mask part need to be reconstructed finally, in order to speed up the training, if the total training round is E, the neural representation loss of all pixels in the image is calculated and sorted after the first round, and then only the pixels with the lowest loss in the next round are trained, and the remaining pixels are not involved in training the neural representation.

[0024] Further, the image restoration dataset distillation method based on uniform image selection and neural representation has the characteristics that, in step 5, given the trained neural representation and the image, the neural representation loss function is calculated, the pixels with the lowest loss in the proportion of r are masked, and the mask coordinates are recorded to obtain the distillation dataset, so that the final pixel proportion is r.

[0025] Further, the image restoration dataset distillation method based on uniform image selection and neural representation has the characteristics that, in step 6, based on the mask coordinates and the neural representation, the pixel value of the pixel located at the mask coordinate (w, h) in the distillation dataset is replaced by the pixel value reconstructed based on the neural representation to reconstruct the distillation dataset, which is used for evaluation of the distillation algorithm.

[0026] Further, the image restoration dataset distillation method based on uniform image selection and neural representation has the characteristics that, in step 7, the restoration model is trained in the distillation dataset, and the peak signal-to-noise ratio and the structural similarity index are calculated as the evaluation indexes of the distillation algorithm.

[0027] Compared with the prior art, the present application has the following advantages:

[0028] (1) The uniform selection method proposed in the present application can simultaneously consider images with different restoration difficulties in the original dataset, ensuring the diversity of the distillation dataset, and greatly compressing the original dataset while maintaining the distribution of the original dataset for the restoration task as much as possible.

[0029] (2) The mask method based on the neural representation proposed in the present application further compresses the dataset at the pixel level, and the final distillation data is obtained by reconstructing the original image based on the neural representation in the deployment stage. Since it is operated at the pixel level, this method is applicable to image restoration datasets of any resolution. BRIEF DESCRIPTION OF DRAWINGS

[0030] ​​The features and advantages of the present application will be appreciated upon reference to the following detailed description and drawings, in which:

[0031] Figure 1 A uniform image selection and neural representation based image restoration dataset distillation method flowchart of the present application

[0032] Figure 2 A uniform image selection and neural representation based image restoration dataset distillation method in different restoration tasks Distillation Result Visualization Chart of the present application, where (ai) is a low-quality distillation image with high peak signal-to-noise ratio, (bi) is a real distillation image with high peak signal-to-noise ratio, (ci) is a low-quality distillation image with medium peak signal-to-noise ratio, (di) is a real distillation image with medium peak signal-to-noise ratio, (ei) is a low-quality distillation image with low peak signal-to-noise ratio, (fi) is a real distillation image with low peak signal-to-noise ratio, where i = 1, 2, 3, 4, 5, 6, respectively representing raindrop removal, rain and fog joint removal, snow removal, fog removal, defogging and denoising.

[0033] Figure 3 A uniform image selection and neural representation based image restoration dataset distillation method in different restoration tasks Comparison result chart with other distillation methods of the present application, where (ai) is a low-quality image, (bi) is the result on the test set after selecting samples with high peak signal-to-noise ratio as the distillation dataset and training on it, (ci) is the result on the test set after selecting samples with low peak signal-to-noise ratio as the distillation dataset and training on it, (di) is the result on the test set after training on the distillation dataset of the proposed method, (ei) is a real image, where i = 1, 2, 3, 4, 5, 6, respectively representing raindrop removal, rain and fog joint removal, snow removal, fog removal, defogging and denoising. DETAILED DESCRIPTION

[0034] The technical solutions in the embodiments of the present application will be described clearly and completely below with reference to the drawings in the embodiments of the present application. Obviously, the described embodiments are part of the embodiments of the present application, rather than all the embodiments. Based on the embodiments in the present application, all other embodiments obtained by those skilled in the art without creative labor should fall within the scope of protection of the present application.

[0035] Reference Figure 1 A uniform image selection and neural representation based image restoration dataset distillation method flowchart, which details the dataset distillation method of the present application.

[0036] Step 1: Train a deep learning model in the original image dataset and specify the final dataset compression ratio r. The input and output of the deep learning model are low-quality images and real images matched with them, respectively. The network architecture includes an encoder and a decoder, which are used to extract low-quality image features and generate real images, respectively. The optimization objective includes a restoration loss function and a perception loss function. The restoration loss function L r is:

[0037]

[0038] where x and y represent low-quality images and real images, respectively, and f θ (.) represents the restoration model with parameters θ. The perception loss function L p is:

[0039]

[0040] where is a pre-trained VGG network.

[0041] Step 2: After generating the restoration images of the images in the training set based on the restoration network trained in Step 1, calculate the peak signal-to-noise ratio between the restoration images and the corresponding real images, and sort them from small to large according to the peak signal-to-noise ratio values. The calculation method of the peak signal-to-noise ratio is:

[0042]

[0043] where w and h represent the width and height of the training image, respectively. The higher the peak signal-to-noise ratio value, the better the image restoration effect, indicating that the training sample is an easy sample for the restoration task. Conversely, the lower the value, the worse the restoration effect, indicating that the sample has a high restoration difficulty and is greatly affected by noise.

[0044] Step 3: After removing the images with the highest and lowest peak signal-to-noise ratios in a proportion of t, select image pairs uniformly according to the peak signal-to-noise ratio. If the low-quality-real image sequence sorted from small to large peak signal-to-noise ratio is {x1,y1},{x2,y2},…{x n ,y n}, then the selected images are:

[0045]

[0046] where n represents the number of images in the original dataset, r1 is the proportion of selected images in the first stage, {x j ,y j} represents the low-quality image and real image pair ranked j from small to large peak signal-to-noise ratio, t is the proportion of images excluded from the highest and lowest peak signal-to-noise ratios, and satisfies 0<t<1, Floor operator. This step filters the original dataset from the image level, aiming to distill more diverse samples in the dataset.

[0047] Step 4: Train the neural representation for each selected image respectively, i.e., the input of the neural network is the coordinate, and the output is the RGB pixel value of the corresponding coordinate. As another storage form of images, the required storage space and the parameters of the neural representation network itself are generally much smaller than the image itself. For an image with a width of w and a height of H, the neural representation loss function of the image is:

[0048]

[0049] where, is the output of the neural representation with parameter θ1 for given coordinate (w, h), and x(w, h) represents the RGB pixel value of the image coordinate (w, h). Both low-quality images and real images in step 3 need to participate in training. Since only the pixels of the mask part need to be reconstructed in the end, in order to speed up the training, if the total training rounds are E, then after rounds, the neural representation loss is calculated for all pixels in the image and sorted, and then in the next rounds, only the pixels with the lowest loss ratio of are trained, while the remaining pixels do not participate in the training of the neural representation.

[0050] Step 5: For each selected image, reconstruct the mask based on the neural representation to the pixel with the smallest loss ratio of , and record the mask coordinates as the final distillation dataset. This step further compresses the dataset from the pixel level. The pixels with the smallest mask reconstruction loss are selected to minimize the difference between the reconstructed pixels and the original image. Since the neural representation-based method directly operates at the pixel level, this method is suitable for image restoration datasets of different resolutions. In addition, the multi-weather restoration dataset includes raindrop removal, rain and fog joint removal, and snow removal data. Since the number of raindrop removal data is an order of magnitude lower than that of the other two tasks, in order to balance the number of images of different tasks, only the neural representation training and masking are performed on the raindrop removal data, and the other two tasks do not participate in the neural representation training.

[0051] Step 6: Based on the mask coordinates, the distillation dataset is reconstructed into a restoration dataset based on the neural representation, i.e., replacing the pixels in the distillation dataset located at the mask coordinates (w, h) with the pixel values reconstructed based on the neural representation Step 7: Reconstruct the distillation dataset to generate a restoration dataset for evaluation of the distillation algorithm. This step is used to generate distillation data that can be evaluated.

[0052] Step 7: Calculate the peak signal-to-noise ratio and structural similarity in the test set after training the deep learning model on the reconstructed distillation data to evaluate the performance of the distillation algorithm. To ensure that each image participates in the same number of iterations, the number of training rounds for the distillation dataset will be reduced by the same proportion as the dataset compression ratio compared to the number of rounds trained on the full dataset.

[0053] Reference Figure 2 Figure 1 is a visualization of the distillation results of the proposed image restoration dataset distillation method based on uniform image selection and neural representation in different restoration tasks. (ai) is a low-quality distillation image with high peak signal-to-noise ratio, (bi) is a real distillation image with high peak signal-to-noise ratio, (ci) is a low-quality distillation image with medium peak signal-to-noise ratio, (di) is a real distillation image with medium peak signal-to-noise ratio, (ei) is a low-quality distillation image with low peak signal-to-noise ratio, and (fi) is a real distillation image with low peak signal-to-noise ratio, where i = 1, 2, 3, 4, 5, 6, representing raindrop removal, rain and fog removal, snow removal, defogging, deblurring, and denoising, respectively. As shown in Figure 2 the higher the peak signal-to-noise ratio, the less the image is affected by rain, fog, snow, blur, and noise, while the lower the peak signal-to-noise ratio, the greater the image is affected by rain, fog, snow, blur, and noise. Images with different levels of noise interference ensure the diversity of the distillation dataset. In addition, the mask based on neural representation usually masks the smooth part of the image, such as the sky area.

[0054] Reference Figure 3 Figure 2 is a comparison of the proposed image restoration dataset distillation method based on uniform image selection and neural representation with other distillation methods in different restoration tasks. (ai) is a low-quality image, (bi) is the result on the test set after selecting samples with high peak signal-to-noise ratio as the distillation dataset and training on it, (ci) is the result on the test set after selecting samples with low peak signal-to-noise ratio as the distillation dataset and training on it, (di) is the result on the test set after training on the distillation dataset proposed in the method, and (ei) is a real image, where i = 1, 2, 3, 4, 5, 6, representing raindrop removal, rain and fog removal, snow removal, defogging, deblurring, and denoising, respectively. The image selection method of (bi) can be represented as:

[0055]

[0056] where n represents the number of images in the original dataset, r is the selection image ratio, {x i ,y i} represents the i-th ranked low-quality image and real image pair with peak signal-to-noise ratio from small to large. The image selection method of (ci) can be represented as:

[0057]

[0058] As Figure 3 shown, compared with other image selection methods, the proposed image restoration dataset distillation method based on uniform image selection and neural representation can achieve better restoration performance.

Claims

1. A data distillation method for image restoration based on uniform image selection and neural representation, characterized in that, Includes the following steps: Step 1: Train a deep learning model on the image restoration dataset, and give the final dataset compression ratio r, where 0 <r<1; Step 2: After inputting low-quality samples from each training set into the model, calculate the peak signal-to-noise ratio between the restored results and the real samples, and sort them from smallest to largest; Step 3: After removing images whose peak signal-to-noise ratio (PSNR) has the lowest to highest ratio (t), uniformly select images according to the compression ratio (r1), where 0 <r<r1<1; Step 4: Train the corresponding neural representation for each selected image. That is, the input of the neural network is coordinates, and the output is the RGB pixel value of the image corresponding to the coordinates. Step 5: For each selected image, the mask is used to reconstruct the image based on the neural representation, minimizing the image loss. The proportion of pixels is recorded, and the mask coordinates are recorded as the final distillation dataset; Step 6: Based on the mask coordinates, the dataset is distilled and the neural representation is used to reconstruct the restored dataset; Step 7: After training a deep learning model on the reconstructed data, calculate the peak signal-to-noise ratio and structural similarity on the test set to evaluate the performance of the distillation algorithm.

2. The image restoration dataset distillation method based on uniform image selection and neural representation according to claim 1, characterized in that: In step 3, the selected image set is uniformly selected from the original dataset according to the peak signal-to-noise ratio (PSNR) value, which can be specifically expressed as: Where n represents the number of images in the original dataset, r1 is the proportion of images selected in the first stage, and {x j ,y j } represents the low-quality image and real image pairs ranked j by peak signal-to-noise ratio (PSNR) from smallest to largest, and t is the proportion of images with the highest and lowest PNR excluded, satisfying 0. <t<1, This represents the floor operator.

3. The image restoration dataset distillation method based on uniform image selection and neural representation according to claim 1, characterized in that: In step 4, for an image with width W and height H, the neural representation loss function is: in, Let x(w,h) represent the output of the neural representation with parameters θ1, given coordinates (w,h), where x(w,h) represents the RGB pixel value at image coordinates (w,h). To accelerate training on low-quality and real images, the neural representation is trained as follows: if the number of training epochs is E, then in... After each iteration, the neural representation loss is calculated and sorted for all pixels in the image, and then in the next iteration... The proportion of those with the lowest losses in the round is Pixel-trained neural representation.

4. The image restoration dataset distillation method based on uniform image selection and neural representation according to claim 1, characterized in that: In step 5, given the trained neural representation and image, the neural representation loss function is calculated, and the proportion with the lowest loss is... The pixels are masked and the mask coordinates are recorded to obtain the distillation dataset.

5. The image restoration dataset distillation method based on uniform image selection and neural representation according to claim 1, characterized in that: In step 6, based on the mask coordinates and neural representation, the pixels located at mask coordinates (w, h) in the distillation dataset are replaced with the pixel values ​​reconstructed based on the neural representation. Reconstruct the distillation dataset for evaluating distillation algorithms.

Images