A low-illumination image enhancement method based on convolutional neural network and wavelet transform

Abstract

The application belongs to the technical field of computer image processing, and particularly relates to a low-illumination image enhancement method based on a convolutional neural network and a wavelet transform. The application comprises: (1) constructing a network SWANet for low-illumination image enhancement, comprising a network MIANet for adjusting the light distribution in the first stage and a network WNENet for removing residual noise and perfecting details in the second stage, wherein a multi-scale attention fusion module based on an attention mechanism and multi-scale features, a channel attention module based on excitation and compression, and a discrete wavelet loss for balancing the spatial domain and frequency features are used. (2) inputting a low-illumination image into the MIANet to obtain a preliminarily enhanced image, and then inputting the preliminarily enhanced image into the WNENet to obtain a final enhanced image. Through the coarse-to-fine enhancement process of the SWANet, the application can effectively adjust the brightness and contrast of different regions, remove noise in the image, and improve the visual experience.

Description

Technical Field

[0001] This invention belongs to the field of computer image processing technology, and specifically relates to a method for enhancing low-light images. Background Technology

[0002] Due to environmental and equipment limitations, photos taken in low-light conditions often suffer from insufficient brightness, low contrast, color distortion, and loss of detail. This not only reduces the aesthetic quality of the image and affects the viewer's visual experience, but also hinders high-level visual tasks such as object recognition and image segmentation. Therefore, low-light image enhancement has become an important and highly anticipated task.

[0003] Low-light image enhancement is primarily implemented in two ways: hardware and software. Hardware implementation refers to improving the acquisition capabilities of hardware devices such as sensors, enabling them to capture clear images under low-light conditions, such as by extending exposure time or using infrared devices for image acquisition. While hardware implementation can overcome the adverse effects of low light to some extent, it is more expensive. Therefore, in practice, software implementation, which is less expensive and more reusable, is often used. Software implementation can be further divided into two types: traditional methods and deep learning methods.

[0004] Traditional methods include enhancement methods based on histogram equalization, gamma correction, and Retinex theory. Histogram equalization and gamma correction methods can improve image visual quality to some extent, but they are not flexible enough and can sometimes disrupt the contrast relationship between adjacent areas, resulting in overexposure. Retinex-based methods ignore noise factors, thus amplifying noise and causing color distortion. Furthermore, the complexity of the optimization process and the long runtime are also drawbacks of these methods.

[0005] With the development of deep learning technology, deep learning-based low-light image enhancement techniques have achieved a leap forward with remarkable results. However, existing methods cannot achieve a balance between adjusting global illumination and local contrast, making it difficult to capture complex noise distributions. Summary of the Invention

[0006] The purpose of this invention is to provide a low-light image enhancement method based on convolutional neural networks and wavelet transform that can obtain high-quality images and excellent visual effects.

[0007] The low-light image enhancement method based on convolutional neural networks and wavelet transform provided by this invention includes the following specific steps:

[0008] (i) Construct a network for low-light image enhancement, denoted as SWANet; the SWANet network includes a network for adjusting the illumination distribution in the first stage (denoted as MIANet) and a network for removing residual noise and improving details in the second stage (denoted as WNENet); among them, a multi-scale attention fusion module based on attention mechanism and multi-scale features, a channel attention module based on excitation and compression, and a discrete wavelet loss for balancing spatial domain and frequency features are used.

[0009] (ii) In the first stage, the low-light image is input into MIANet to obtain the preliminarily enhanced image; in the second stage, the preliminarily enhanced image is input into WNENet to obtain the final enhanced image.

[0010] This can be formally expressed as follows:

[0011] I1 = MIANet(I)

[0012] I2 = WNENet(I1)

[0013] Where I, I1, and I2 represent the input image, the initially enhanced image, and the final enhanced image, respectively.

[0014] In step (1), the MIANet network is designed based on the UNet structure and consists of two parts: a multi-scale feature extraction module and a multi-scale attention fusion module.

[0015] In the multi-scale feature extraction module, the input image is downsampled sequentially using bilinear interpolation at 2x, 4x, and 8x, resulting in images of the same size as the original image. Figure 1 The original image is taken as half, 1 / 4, and 1 / 8 of the original image. Then, a weight-shared 3x3 convolution is performed between the original image and these three images at different scales to obtain features at four scales. The original image is then processed through an NAFBlock. [1] And a 1x1 convolution, the output is fed into a pixel unshuffle for downsampling. [7] The operation is performed with a 3x3 convolution to obtain downsampled features down1 and encoded features E1; the features at the second and third scales are downsampled by themselves and the result of downsampling the features at the previous scale. i-1 The components are concatenated and then sequentially fed into an NAFBlock and a 1x1 convolution. The outputs are then fed into the Pixel Unshuffle function. [7] Downsampling and a 3x3 convolution are used to obtain downsampled features. i and encoding feature E i (i = 2, 3); The feature at the 4th scale concatenates D3 with itself, and after one NAFBlock and a 1x1 convolution, the encoded feature E4 is obtained. Encoded feature E4 is then processed by Pixel Shuffle.[7] Upsample the sample and concatenate it with E3, then feed it into an NAFBlock and a 3x3 convolution to obtain the decoded feature D2; repeat the above operation on the second scale to obtain the decoded feature D2. Concatenate E1 with the upsampled D2 and feed it into an NAFBlock and a 3x3 convolution to obtain the decoded feature D1.

[0016] After obtaining feature information at multiple scales, they are fed into the multi-scale attention fusion module to obtain the initially enhanced image I1.

[0017] The workflow of the multi-scale attention fusion module is as follows: First, the four decoded input features are upsampled to the same size as the original image using bilinear interpolation and then summed. Then, the result is passed through the convolutional block attention module. [5] Next, features are extracted from the original image using a 3x3 convolution and summed. Then, the image is passed through a Conv-Conv-ReLU-Conv-ReLU-Conv block containing skip connections, and the output is summed with the original image to obtain the preliminarily enhanced image I1.

[0018] In step (one), the specific working process of the WNENet network is as follows:

[0019] For the input image I1, it is first segmented into four frequency sub-bands using discrete wavelet transform, and then fed into four 3x3 convolutions. The four outputs are then concatenated and fed into a 3x3 convolution, ReLU, and channel attention module. Next, a 1x1 convolution is used to adjust the number of channels, and the features are transformed into the enhanced image I2 using inverse discrete wavelet transform. During this training process, discrete wavelet loss is used for supervised training.

[0020] The specific workflow of the channel attention module is as follows:

[0021] The input features are first processed through a Conv-ReLU-Conv block, and then through an inverse SE block using an "excitation-compression" operation to adjust the feature weights. The inverse SE block and the SE... [2] The block is similar, but instead of reducing the channel dimensionality by 8 times in the SE block, the channel dimensionality is increased by 128 times. Then it is added to the residual, and the channel output is adjusted using a 1x1 convolution.

[0022] The discrete wavelet loss method is employed, and the specific process is as follows:

[0023] Given the enhancement result Given a standard illumination image Y corresponding to a low-light image, the discrete wavelet loss method first transforms both images to the frequency domain using discrete wavelet transform, then uses Charbonnier loss... [6] The loss function supervises the training process. This process can be expressed as:

[0024]

[0025] Among them, f DWT This represents the discrete wavelet transform operation, with ε being 0.001.

[0026] In this invention:

[0027] Training the MIANet network includes:

[0028] (a) Input a low-light image I to obtain a preliminarily enhanced image I1;

[0029] (b) Set training parameters, including batch size, image patch size, learning rate, learning rate adjustment strategy, and optimizer;

[0030] (c) Begin training;

[0031] (d) Adjust the hyperparameters in the model multiple times to obtain the model with the best performance.

[0032] Training the WNENet network includes:

[0033] (1) Use MIANet to process low-light images to obtain preliminary enhanced images, divide them into training and test sets according to the proportion, and perform image preprocessing.

[0034] (2) Train the WNENet network;

[0035] During training, discrete wavelet loss is used to balance spatial and frequency domain features, thereby facilitating the training process.

[0036] (3) Set training parameters, including batch size, image patch size, learning rate, learning rate adjustment strategy and optimizer;

[0037] (4) Adjust the hyperparameters in the model multiple times to obtain the model with the best performance.

[0038] Compared with existing technologies, this invention can adjust uneven lighting in images, remove noise from images, obtain higher image quality indicators and achieve better visual effects. It can be applied to scenarios such as video surveillance and has good application value. Attached Figure Description

[0039] Figure 1 This is a diagram of the SWANet network structure and channel attention block structure used in this invention.

[0040] Figure 2 This is the multi-scale attention fusion module used in this invention.

[0041] Figure 3 The results show the visual comparison with the state-of-the-art method on the real low-light dataset VV. Detailed Implementation

[0042] The network structure used in this invention is as follows: Figure 1 As shown. The input low-light image is first initially adjusted using MIANet in the first stage, and then further refined for image details using WNENet. Channel attention blocks are in Figure 1 In the bottom right corner, the multi-scale attention fusion module is as follows: Figure 2 As shown.

[0043] The embodiments of the present invention include the following steps:

[0044] (1) Dataset preparation

[0045] This invention in LOL [3] and MIT Adobe FiveK [4] Training and testing are performed on the dataset. LOL is the first dataset captured in a real-world scene, containing 500 image pairs, of which 485 pairs are used for training and 15 pairs for testing. The MIT / Adobe FiveK dataset was collected by Bychkovsky et al. at MIT. It was initially used for global tone adjustment of images and later for low-light image enhancement, containing 5000 RAW format images. In this example, the original images are converted from RAW to JPG, and the first 4500 image pairs are used for training, while the last 500 pairs are used for testing.

[0046] (2) Data preprocessing

[0047] Since the number of training images is limited, data augmentation operations are introduced in this embodiment of the invention to increase the amount of training data and ensure that the model is fully trained. First, image patches are randomly cropped from the input image, and then randomly flipped horizontally, vertically, and rotated.

[0048] (3) Parameter settings

[0049] This embodiment of the invention uses a GTX 3090 GPU for training, and the model's hyperparameters are set as follows: batch size of 4; input image size of 256x256; optimizer of AdamW; and in the first stage, the initial learning rate is set to 5×10. -5 The learning rate was reduced to 1×10 using a cosine annealing strategy. -5 In the second phase, the initial learning rate was set to 1×10. -4 The learning rate was reduced to 1×10 using a cosine annealing strategy. -5 Each stage requires training the network for 300 epochs.

[0050] (4) Experimental Results

[0051] Figure 3 The test results on the VV dataset are shown. (a) is the input image, (b)-(k) are other state-of-the-art (SOTA) methods, and (i) is the method proposed in this invention. It can be seen that LIME simply brightens the entire image, resulting in significant noise. The other enhancement methods produce underexposed or overexposed results. In contrast, this embodiment brightens the overall image while balancing local contrast, without significant noise, resulting in better image quality and a superior visual effect.

[0052] Table 1 presents the comparison results of the present invention with 11 other enhancement methods on the LOL and MIT datasets. As can be seen, compared with other methods, the embodiments of the present invention achieved the highest performance metrics on both datasets, confirming the effectiveness and superiority of the present method.

[0053] Table 1

[0054]

[0055]

[0056] References

[0057] [1]Chen L,Chu X,Zhang X,et al.Simple baselines for image restoration[C] / / Proceedings of the European Conference on Computer Vision.2022:17-33.

[0058] [2]Hu J, Shen L, Sun G. Squeeze-and-excitation networks[C] / / Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition.2018:7132-7141.

[0059] [3]Wei C, Wang W, Yang W, et al. Deep retinex decomposition for low-lightenhancement[J]. BMVC, 2018.

[0060] [4]Bychkovsky V,Paris S,Chan E,et al.Learning photographic globaltonal adjustment with adatabase of input / output image pairs[C] / / Proceedingsof the IEEE Conference on Computer Vision and Pattern Recognition.2011:97-104.

[0061] [5]Woo S,Park J,Lee J Y,et al.Cbam:Convolutional block attentionmodule[C] / / Proceedings of the European Conference on Computer Vision.2018:3-19.

[0062] [6]Lai W S,Huang J B,Ahuja N,et al.Fast and accurate image super-resolution with deep laplacian pyramid networks[J].IEEE Transactions onPattern Analysis and Machine Intelligence,2018,41(11):2599-2613.

[0063] [7]Shi W,Caballero J,Huszár F,et al.Real-time single image and videosuper-resolution using an efficient sub-pixel convolutional neural network[C] / / Proceedings of the IEEE Conference on Computer Vision and PatternRecognition.2016:1874-1883。

Claims

1. A low-illumination image enhancement method based on a convolutional neural network and a wavelet transform, characterized in that, The specific steps are: (I) Construct a network for low-illumination image enhancement, denoted as SWANet; the SWANet network includes a network for adjusting the illumination distribution in the first stage, denoted as MIANet, and a network for removing residual noise and perfecting details in the second stage, denoted as WNENet; wherein, a multi-scale feature extraction module and a multi-scale attention fusion module based on attention mechanism, a channel attention module based on excitation and compression, and a discrete wavelet loss for balancing spatial domain and frequency features are used; (II) In the first stage, the low-illumination image is input into the MIANet to obtain a preliminary enhanced image; in the second stage, the preliminary enhanced image is input into the WNENet to obtain a final enhanced image; the formalized expression is as follows: I1=MIANet(I) I2=WNENet(I1) Wherein, I, I1, I2 respectively represent the input image, the preliminary enhanced image, and the final enhanced image; The MIANet network in step (I) is designed based on the UNet structure, including a multi-scale feature extraction module and a multi-scale attention fusion module; In the multi-scale feature extraction module part, the input image is sequentially down-sampled by bilinear interpolation by 2x, 4x, 8x, and three images with sizes of 1 / 2, 1 / 4, 1 / 8 of the original image are obtained, then the original image and the three images with different scales are subjected to a weight-shared 3x3 convolution to obtain four scale features; then the original image is subjected to a NAFBlock and a 1x1 convolution, and the outputs are respectively sent to a PixelUnshuffle operation for down-sampling and a 3x3 convolution to obtain down-sampling feature down1 and encoding feature E1; the features of the second and third scales are concatenated with the down-sampling results of the features of the previous scales down i-1 , and are sequentially sent to a NAFBlock and a 1x1 convolution, and the outputs are respectively sent to a PixelUnshuffle down-sampling and a 3x3 convolution to obtain down-sampling feature down i and encoding feature E i (i = 2, 3); the fourth scale feature concatenates D3 with itself, and after being subjected to a NAFBlock and a 1x1 convolution, the encoding feature E4 is obtained; the encoding feature E4 is subjected to a PixelShuffle up-sampling and is concatenated with E3, and then is sent to a NAFBlock and a 3x3 convolution to obtain decoding feature D2; the above operation is repeated at the second scale to obtain decoding feature D2; E1 is concatenated with the up-sampled D2 and is sent to a NAFBlock and a 3x3 convolution to obtain decoding feature D1; After obtaining the feature information of multiple scales, they are sent to the multi-scale attention fusion module to obtain the preliminary enhanced image I1; The specific working process of the WNENet network in step (I) is as follows: For the input image I1, first use discrete wavelet transform to cut it into 4 frequency subbands, respectively input into 4 3x3 convolutions, then concatenate the 4 outputs, and input into 3x3 convolution, ReLU, channel attention module, then use 1x1 convolution to adjust the channel number and use inverse discrete wavelet transform to convert the feature into an enhanced picture I2; in this training process, the discrete wavelet loss is used for supervised training; The specific working process of the channel attention module is as follows: First, the input feature is passed through a Conv-ReLU-Conv block, and then an inverse SE block using the "excitation-compression" operation is used to adjust the feature proportion; the inverse SE block is to replace the 8-fold channel dimension reduction operation in the SE block with a 128-fold channel dimension increase operation; then add the residual, and use 1x1 convolution to adjust the output of the channel.

2. The low-light image enhancement method based on convolutional neural network and wavelet transform according to claim 1, characterized in that, The working process of the multi-scale attention fusion module is as follows: First, the 4 decoded features are upsampled to the same size as the original image using bilinear interpolation, and then added together, then passed through a convolution block attention module, and then use 3x3 convolution to extract features from the original image, and add them together; then pass through a Conv-Conv-ReLU-Conv-ReLU-Conv block containing a skip connection, and add the output result to the original image to obtain the preliminary enhanced image I1. 3.The low-light image enhancement method based on convolutional neural network and wavelet transform according to claim 1, characterized in that, The supervised training using the discrete wavelet loss is as follows: Given the enhancement result And the standard illumination image Y corresponding to the low illumination image pair, the discrete wavelet loss first converts the two images into the frequency domain by discrete wavelet transform, and then supervises the training process by the Charbonnier loss function; the process is expressed as: where f DWT denotes a discrete wavelet transform operation and ε is 0.001.

Images