A method for underwater image enhancement of a lightweight convolutional network

Abstract

The present application relates to a kind of light convolution network's underwater image enhancement method, based on the traditional white balance algorithm in neural network introduces trainable white balance module, balance the color intensity of RGB three channels, can balance color channel.For the problem of low contrast, the present application is based on the traditional histogram stretching algorithm again and introduces multiple space trainable histogram stretching module, in RGB, Lab, HSI three spaces, and the module is extended to multiple color space parallel operation, can improve image contrast.In addition, residual enhancement module is also introduced to restore image details.Compared with the current SOTA model, the best enhancement effect is obtained.The parameter quantity is reduced by more than 98%.Greatly reduce the image enhancement time, can be used to enhance high-resolution image and underwater video.

Description

TECHNICAL FIELD

[0001] The present application belongs to the field of image processing, and relates to a method for underwater image enhancement based on a light-weight convolutional network, which is suitable for image and video enhancement in underwater environments. BACKGROUND

[0002] At present, underwater enhancement techniques can be roughly divided into two categories: traditional algorithms and model algorithms based on deep learning.

[0003] Traditional algorithms can be further divided into algorithms based on physical models and algorithms based on non-physical models. Among them, the algorithms based on physical models mainly estimate the parameters of the underwater physical image model by introducing priori, and the representative priori methods include the red channel prior (Red Channel Prior), the underwater dark channel prior (Underwater Dark Channel Prior), the underwater light attenuation prior (Underwater Light Attenuation Prior), etc. The non-physical model-based method improves the image quality by adjusting the image pixel value, and the representative method includes the unsupervised color correction method (Unsupervised Colour Correction Method) and the image enhancement method based on relative global histogram stretching based on adaptive parameter acquisition (Relative Global Histogram Stretching Based on Adaptive Parameter Acquisition). However, whether it is the physical prior introduced by the physical model-based algorithm or the pixel adjustment strategy adopted by the non-physical model-based algorithm, the application scope of these algorithms is relatively narrow, and it is difficult to be applied to various underwater environments.

[0004] Therefore, with the rapid development of deep learning networks in the field of image processing, many powerful models for underwater image enhancement have emerged. For example, the Waternet, Ucolor and UWCNN models based on convolutional neural networks (CNN), and the FUnIE-GAN and UWGAN models based on generative adversarial networks (GAN). These models benefit from the powerful learning ability of neural networks and can enhance images to a very high quality level. However, most of these methods refer to the model architecture applied to other visual tasks, ignoring the prior knowledge of traditional underwater enhancement algorithms, resulting in thousands of parameters in the model, leading to excessive training time and processing picture time, and easy overfitting on small data sets.

[0005] Compared with ground images, underwater images are affected by scattering and absorption of light, affecting the reliability and quality of image processing. Underwater images mainly have low contrast and color imbalance problems. For the problem of color imbalance, the method based on deep learning in the application introduces priori of traditional algorithm, and the image enhancement effect exceeds the state-of-the-art (SOTA) model, and compared with the existing SOTA model, the parameter quantity is reduced by more than 98%. SUMMARY

[0006] TECHNICAL PROBLEM TO BE SOLVED

[0007] In order to avoid the shortcomings of the prior art, the application provides a method for underwater image enhancement of a lightweight convolutional network, which introduces a trainable white balance module to balance the color intensity of the RGB three channels. For the problem of low contrast, the application introduces a multi-space trainable histogram stretching module to increase the image contrast in RGB, Lab and HSI three spaces in parallel. In addition, a residual enhancement module is also introduced to restore image details.

[0008] TECHNICAL SCHEME

[0009] A method for underwater image enhancement of a lightweight convolutional network, characterized by the following steps:

[0010] Step 1: Trainable white balance module

[0011] Step a1: Calculate the mean and maximum of the RGB three channels of the original image

[0012] Step a2: Copy the original image into two copies, process them with a convolutional layer, and then fuse the values of M and A respectively to obtain two pictures balanced according to different reference values; C and A C

[0013] Step a3: Feature fusion is performed on the two pictures, and the feature fusion steps are to perform convolutional layer processing on the two balanced pictures respectively, and then add them pixel by pixel;

[0014] Step a4: The enhanced image obtained by inputting the original image into the residual enhancement module is added to the picture obtained in step a3 pixel by pixel;

[0015] The residual enhancement module includes a convolutional layer and an activation function;

[0016] Step 2: Trainable histogram stretching module

[0017] Step b1: Convert the picture to HSI and Lab spaces respectively;

[0018] Step b2: Perform trainable histogram stretching in RGB, HSI and Lab spaces;​

[0019] Step b3: convert the pictures on HSI and Lab back to RGB color space to get three enhanced pictures;

[0020] Step b4: feature fusion is performed on the three pictures, and the feature fusion step is to perform convolution layer processing on the three stretched pictures respectively, and then add them pixel by pixel;

[0021] Step 3: add the picture obtained in step a4 to the picture obtained in step b4 pixel by pixel through the residual enhancement module, so that the underwater image of the lightweight convolution network is enhanced.

[0022] The training of the trainable white balance module and the trainable histogram stretching module uses the UIEB dataset.

[0023] The batchsize is set to 16, the epoch is set to 200, the optimization algorithm is Adam, and the learning rate is 1e-3 during training.

[0024] The UIEB dataset contains 890 pairs of datasets, of which 800 pairs are training sets and 90 pairs are test sets. The model parameters are trained on the training set.

[0025] Beneficial effects

[0026] The method for enhancing underwater images of a lightweight convolution network proposed in the application introduces a trainable white balance module based on a traditional white balance algorithm in a neural network, balances the color intensity of the RGB three channels, and can balance the color channels. In view of the problem of low contrast, the application introduces a multi-space trainable histogram stretching module based on a traditional histogram stretching algorithm, in RGB, Lab, HSI three spaces, and extends the module to multiple color spaces and parallel operation, which can improve the contrast of the image. In addition, a residual enhancement module is introduced to restore the image details.

[0027] Now the mainstream underwater enhancement model is based on a deep learning model, but these models often have a large number of parameters, resulting in long training time, difficulty in fitting, and being not conducive to real-time enhancement processing of underwater video and high-resolution underwater image enhancement. The application introduces the traditional algorithms of white balance and histogram stretching, so only a small number of parameters are needed to achieve good results. As shown in Table 1, compared with other SOTA deep learning algorithms, the parameter quantity is reduced by more than 98%, and the number of floating point operations per second is reduced by more than 93%. Therefore, the application greatly improves the picture enhancement speed, reduces the memory requirement of the device, and can process high-resolution underwater images.

[0028] Now the mainstream underwater enhancement model is based on deep learning model, but these models often have excessive parameter quantity, resulting in long training time, difficult to fit, and not conducive to real-time enhancement processing of underwater video and high-resolution underwater image enhancement. The present application because of the introduction of the traditional algorithm of white balance and histogram stretching priori, so only a small amount of parameters can achieve good results. As shown in Table 1, compared with other SOTA deep learning algorithms, the parameter quantity is reduced by more than 98%, and the number of floating point operations per second is reduced by more than 93%. Therefore, the present application greatly improves the picture enhancement speed, reduces the memory requirement of the device, and can process high-resolution underwater images.

[0029] Table 1

[0030] parameters floating-point operations per second (GFLOPs) the invention 0.001M 0.05 water network 1.09M 142.91 UWCNN 0.04M 0.79 FUnIE-GAN 7.02M 0.73 Ucolor 148.77M 2805.34

[0031] The present application proposes a method for underwater image enhancement of a lightweight convolutional network, and the beneficial effects are as follows:

[0032] 1. The method disclosed in the present application achieves the best enhancement effect on the image enhancement data set (UIEB) compared with the current SOTA model.

[0033] 2. Compared with the mainstream deep learning model applied to underwater enhancement, the parameter quantity of the method disclosed in the present application is reduced by more than 98%.

[0034] 3. The method disclosed in the present application greatly reduces the image enhancement time and can be used to enhance high-resolution images and underwater videos. BRIEF DESCRIPTION OF DRAWINGS

[0035] figure 1 is a model overall framework diagram;

[0036] figure 2 Five underwater original images are given, and the numbers in the upper right corner are PSNR / SSIM. The larger the value, the closer to the label picture, and the better the visual effect.

[0037] figures 3 to 7 The results of four SOTA models and the method of the present application are given figure 2 (a) to figure 2 (e), figures 3 to 7 (a) is the result of RGHS; (b) is the result of ULAP; (c) is the result of FUnIE-GAN; (d) is the result of Ucolor; (e) is the result of the present application; all pictures have numbers in the upper right corner PSNR / SSIM. The larger the value, the closer to the label picture, and the better the visual effect. DETAILED DESCRIPTION

[0038] The application will be further described in connection with the embodiments and drawings as follows:

[0039] A method for underwater image enhancement of a light-weight convolutional network, the main modules and corresponding steps are as follows:

[0040] Module 1: trainable white balance module

[0041] In order to obtain high-quality images, the color channel values of RGB need to be equal, wherein the specific reference values of the trainable white balance module are the mean and maximum values of the RGB color channels. The trainable white balance module can be divided into three steps.

[0042] Step 1: calculate the mean and maximum values of the three channels of the image RGB:

[0043]

[0044]

[0045] wherein I C (i, j) represents the pixel values of the red, green and blue components of the RGB image, C represents the channel corresponding to the RGB image, MN represents the size of the image, M C and A C respectively represent the maximum and mean values of each color component of the obtained RGB image.

[0046] Step 2: the original image is copied into two copies, processed using a convolutional layer, and then fused with the values of M C and A C , to obtain two pictures balanced according to different reference values.

[0047] Step 3: feature fusion is performed on the two pictures.

[0048] Module 2: trainable histogram stretching module

[0049] In order to improve the contrast of the image, a trainable histogram stretching module is introduced, and enhancement is performed in multiple color spaces. The trainable histogram stretching module can be divided into four steps.

[0050] Step 1: convert the pictures into HSI and Lab spaces respectively.

[0051] Step 2: perform trainable histogram stretching in RGB, HSI and Lab spaces.

[0052] The formula of histogram stretching can be expressed as:

[0053]

[0054] wherein P0 is the pixel value after contrast correction; P iFor the pixel value of the current image, a is the minimum pixel value existing in the enhanced image, b is the minimum pixel value existing in the enhanced image. C is the minimum pixel value existing in the pre-enhanced image, and d is the maximum pixel value existing in the pre-enhanced image.

[0055] In formula (3), a and b are known values, and c and d are unknown values. The trainable histogram stretching module can train these two values using a convolutional layer.

[0056] Step 3: Convert the pictures on HSI and Lab back to the RGB color space to obtain three enhanced pictures.

[0057] Step 4: Perform feature fusion on the three pictures.

[0058] Module 3: Residual Enhancement Module

[0059] The residual enhancement module is attached to module 1 and module 2, which can be used to improve the fidelity of image information and supplement the details of the picture. It can prevent the phenomenon of over-enhancement and under-enhancement of the picture.

[0060] In specific embodiments:

[0061] In actual use, different parameters can be trained for different underwater scenes to improve the accuracy of the model. The underwater image public dataset UIEB contains various underwater scenes, such as color imbalance, low contrast, and picture fog. For the dataset UIEB, the specific implementation process is as follows:

[0062] 1: Divide the dataset. UIEB contains 890 pairs of datasets, of which 800 pairs are training sets and 90 pairs are test sets. Train the model parameters on the training set. The batchsize is set to 16, the epoch is set to 200, the optimization algorithm is Adam, and the learning rate is 1e-3.

[0063] 2: Select test pictures and input them into the model. Referring to figure 2 , the five pictures are sampled from the original underwater pictures in the UIEB test set. The numbers in the upper right corner of the pictures represent PSNR / SSIM, which represent the content similarity and structural similarity of the pictures and the corresponding label pictures, respectively. The higher the PSNR and SSIM values, the higher the similarity to the label picture, and the better the picture quality. figure 2 The pictures in the picture are affected by the underwater environment to different degrees. We arrange them from left to right according to the size of PSNR / SSIM, where the leftmost picture is affected more by the quality, and the rightmost picture is affected less by the quality.

[0064] The input picture is processed by module 1, i.e. the white balance module for color balance.

[0065] The picture enhanced by module 1 is subjected to module 2, that is, a histogram stretching module is trained to improve contrast.

[0066] Referring to figures 3 to 7 For the fairness of comparison, we select four SOTA models based on different mechanisms to compare with the present application. The four models are a non-physical model based method (Relative Global Histogram Stretching Based on Adaptive Parameter Acquisition, RGHS), a physical model based method (Underwater Light Attenuation Prior, ULAP), a GAN based method (FUnIE-GAN) and a CNN based method Ucolor. We can see that other models either do not enhance enough or introduce undesirable colors due to over-enhancement. In addition, for figure 7 For such high-quality underwater images, other SOTA models have reduced the SSIM value, resulting in negative optimization. However, our model improves the visual effect of underwater images regardless of their quality, and restores the true colors of these underwater images. Compared with the four SOTA models, the present application achieves the highest PSNR and SSIM in the five photos.

[0067] The final image enhancement comparison results are shown in Table 2. Compared with the SOTA model, the present application achieves the highest PSNR and SSIM in the underwater dataset UIEB. The PSNR and SSIM are 24.43 and 0.91, respectively.

[0068] Table 2: Comparison of enhancement results of the present application and SOTA models

[0069] PSNR SSIM RGHS enhancement 19.12 0.83 ULAP enhancement 16.39 0.75 FUnIE-GAN enhancement 17.77 0.76 Ucolor enhancement 21.95 0.88 invention enhancement 23.43 0.91

Claims

1. A method for underwater image enhancement using a lightweight convolutional network, characterized in that... The steps are as follows: Step 1: Train the white balance module Step a1: Calculate the mean and maximum values ​​of the three RGB channels of the original image. Step a2: Duplicate the original image into two copies, process them using a convolutional layer, and then fuse them separately. C and A C The value is used to obtain two images balanced according to different reference values; Step a3: Perform feature fusion on the two images. The steps are to process the two balanced images separately using convolutional layers, and then add them pixel by pixel; Step a4: Input the original image into the enhanced image obtained by the residual enhancement module, and then add it pixel by pixel to the image obtained in step a3; The residual enhancement module includes a convolutional layer and an activation function; Step 2: Training the histogram stretching module Step b1: Convert the image to HSI and Lab color spaces respectively; Step b2: Perform training histogram stretching in RGB, HSI, and Lab spaces; Step b3: Convert the images from HSI and Lab back to the RGB color space to obtain three enhanced images; Step b4: Perform feature fusion on the three images. The feature fusion process involves processing the three stretched images separately with convolutional layers and then adding them pixel by pixel. Step 3: The image obtained in step a4 is processed by the residual enhancement module and the image obtained in step b4 are added pixel by pixel to enhance the underwater image of the lightweight convolutional network.

2. The underwater image enhancement method using a lightweight convolutional network according to claim 1, characterized in that: The training of the white balance module and the histogram stretching module uses the UIEB dataset.

3. The underwater image enhancement method using a lightweight convolutional network according to claim 2, characterized in that: During training, the batch size was set to 16, the epochs to 200, the optimization algorithm to Adam, and the learning rate to 1e-3.

4. The underwater image enhancement method using a lightweight convolutional network according to claim 2, characterized in that: The UIEB dataset contains 890 dataset pairs, including 800 training pairs and 90 test pairs. Model parameters are trained on the training sets.

Images