An underwater image enhancement method based on a dual-channel convolutional neural network

Abstract

The present invention provides an underwater image enhancement method based on a dual-channel convolutional neural network, which relates to the technical field of underwater image enhancement and restoration, and includes the following steps: S1: Establish an image training set; S2: Input an image to be optimized into the shallow convolutional neural network branch to optimize the global structure of the target image, and obtain a globally structure-optimized image; S3: Input another image to be optimized into the deep convolutional neural network branch to restore the detailed information of the target image, and obtain a detailed information-restored image; S4: Input the globally structure-optimized image and the detailed information-restored image into the fine-tuning convolutional neural network, and perform image fusion using tensor splicing operation to output the enhanced target image; S5: Perform the above steps on all images in the image training set until the image enhancement is completed. The present invention can effectively enhance underwater images with low contrast and serious color distortion, while retaining the detailed information of the images, and can better improve the image quality.

Description

Technical Field

[0001] The present invention relates to the technical field of underwater image enhancement and restoration, and in particular, to an underwater image enhancement method based on a dual-channel convolutional neural network. Background Art

[0002] Currently, underwater images, as an important carrier and manifestation of underwater environmental information, are very important in aspects such as underwater environment exploration, development, and protection. However, due to many adverse factors in the marine environment, such as high turbidity, uneven illumination, and complex background, the obtained underwater images often have problems such as low contrast, color distortion, blurred texture, and quality degradation. The acquisition of high-quality underwater images faces great challenges. Therefore, enhancing and restoring low-quality underwater images is one of the effective means to obtain high-quality underwater images.

[0003] Currently, underwater image enhancement and restoration methods can be divided into three categories: non-physical model-based methods, physical model-based methods, and deep learning-based methods. Non-physical methods do not rely on the underwater optical imaging model and improve the visual quality of the image by adjusting the pixel values of the image. Physical model-based methods mathematically model the degradation process of underwater images and obtain clear underwater images by estimating model parameters and inverting the degradation process. Deep learning-based methods construct deep neural networks and are trained with a large number of underwater images and high-quality reference images to obtain clear underwater images.

[0004] However, the existing underwater image enhancement and restoration methods have certain drawbacks: non-physical model-based methods are prone to introducing color deviations and artifacts and may also exacerbate noise because they do not consider the optical characteristics of underwater images; physical model-based methods are affected by the diversity of image processing and quality measurement methods, and underwater image processing methods cannot adapt to the complex diversity of the marine environment; deep learning-based methods have poor model generalization due to the influence of underwater image datasets, and currently, single-channel convolutional neural networks are often designed for only one problem and still need further research. Summary of the Invention

[0005] In view of this, the purpose of the present invention is to propose an underwater image enhancement and restoration method based on a dual-channel convolutional neural network, and at the same time select multiple types of underwater image datasets to solve the technical problems of low image quality and poor model generalization in the existing method of using a single-channel convolutional neural network for underwater image enhancement and restoration.

[0006] The technical means adopted by the present invention are as follows:

[0007] An underwater image enhancement method based on a dual-channel convolutional neural network, the dual-channel convolutional neural network includes a shallow convolutional neural network, a deep convolutional neural network, and a fine-tuning convolutional neural network, and the underwater image enhancement method includes the following steps:

[0008] S1: Establish an image training set, and select an image from the image training set as the target image to be optimized;

[0009] S2: Input a target image to be optimized into the shallow convolutional neural network branch to optimize the global structure of the target image, and obtain a globally structured optimized image; the shallow network contains five convolutional layers, the convolutional kernels are all 3, and after each convolutional layer, a normalization processing function and a LeakyReLU activation function are connected, and residuals are added in the first convolutional block and the last convolutional block;

[0010] S3: Input another target image to be optimized into the deep convolutional neural network branch to restore the detailed information of the target image, and obtain a detailed information restored image; the deep network contains thirteen convolutional layers, the convolutional kernels are all 3, and after each convolutional layer, a normalization processing function and a LeakyReLU activation function are connected, and the attention mechanism module SE Block and the convolutional block are used to jointly process the target image to be optimized;

[0011] S4: Input the globally structured optimized image and the detailed information restored image into the fine-tuning convolutional neural network, perform image fusion using tensor splicing operations, and output the enhanced target image; the fine-tuning convolutional neural network contains six convolutional layers, the convolutional kernels are all 3, and after each convolutional layer, a normalization processing function and a LeakyReLU activation function are connected;

[0012] S5: Perform the above steps on all images in the image training set until the image enhancement is completed.

[0013] Further, S1 includes the following steps:

[0014] Select underwater image datasets in different ocean environments and at different depths as the training set, and use data augmentation to increase the data volume

[0015] Collect real underwater image datasets in different sea areas and at different depths and synthetic underwater image sets in different scenarios and conditions as the training set and the validation set for the dual-channel convolutional neural network;

[0016] Adopt a data augmentation method to double the sample size of the training set and the validation set.

[0017] Further, the data augmentation method includes image flipping, image scaling, image cropping, and image translation.

[0018] Further, the target image to be optimized is a 3-channel RGB image.

[0019] Further, S2 includes the following steps:

[0020] Perform preprocessing on the target image to be optimized to reduce the influence of noise in the image;

[0021] Use four convolutional neural network modules to restore the global structural information of the preprocessed target image to be optimized;

[0022] Adopt a residual operation to connect the second convolutional block and the last convolutional block of the shallow convolutional neural network, enabling the network to extract deeper features and avoiding gradient vanishing or explosion, thereby obtaining a globally structurally optimized image.

[0023] Further, S3 includes the following steps:

[0024] Perform preprocessing on the target image to be optimized to reduce the influence of irrelevant information in the image;

[0025] Use four self-attention modules to further extract information from the preprocessed target image to be optimized;

[0026] Adopt a channel attention mechanism SE Block to fully distinguish between valid and invalid information. The SE Block starts from a single image, extracts image features, and the feature map dimension of the current feature layer U is [C, H, W]; where H represents height, W represents width, and C represents the number of channels; perform average pooling or max pooling on the [H, W] dimension of the feature map, and the size of the pooled feature map is [C, 1, 1], that is, the weight of each channel is obtained; apply the weight to the feature map U[C, H, W], that is, each channel multiplies itself by its respective weight, to obtain a detail information restored image.

[0027] Further, S4 includes the following steps:

[0028] Fuse the two images obtained after being processed by the shallow and deep convolutional neural network modules together as a new input image;

[0029] Use six convolutional blocks to perform color correction on the new input image;

[0030] Output the enhanced underwater image.

[0031] The present invention also provides a storage medium, which includes a stored program. When the program runs, it executes the underwater image enhancement method based on a dual-channel convolutional neural network described in any one of the above.

[0032] The present invention also provides an electronic device, including a memory, a processor, and a computer program stored on the memory and executable on the processor. The processor runs through the computer program to execute the underwater image enhancement method based on a dual-channel convolutional neural network according to any one of the above.

[0033] Compared with the prior art, the present invention has the following advantages:

[0034] The present invention decomposes the underwater image into a global component and a local component by designing deep and shallow convolutional neural network modules. By separately processing the underwater image enhancement problem in terms of structure and details, a better enhancement effect can be obtained.

[0035] The design of fine-tuning the convolutional neural network module in the present invention can further strengthen the coupling between the deep and shallow modules, which enables the proposed method to handle more challenging underwater image enhancement tasks. BRIEF DESCRIPTION OF THE DRAWINGS

[0036] In order to more clearly illustrate the technical solutions in the embodiments of the present invention or the prior art, the following will briefly introduce the drawings required for the description of the embodiments or the prior art. Obviously, the drawings in the following description are some embodiments of the present invention. For those of ordinary skill in the art, without creative efforts, other drawings can be obtained based on these drawings.

[0037] Figure 1 It is a flowchart of the method of the present invention.

[0038] Figure 2 It is a process diagram of underwater image processing of the present invention.

[0039] Figure 3 It is a specific structure diagram of the self-attention module in the deep branch of the present invention.

[0040] Figure 4 It is a flowchart of the working process of the channel attention mechanism of the present invention. DETAILED DESCRIPTION OF THE EMBODIMENTS

[0041] In order to enable those skilled in the art to better understand the solution of the present invention, the following will clearly and completely describe the technical solutions in the embodiments of the present invention with reference to the drawings in the embodiments of the present invention. Obviously, the described embodiments are only some embodiments of the present invention, rather than all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those of ordinary skill in the art without creative efforts shall fall within the protection scope of the present invention.

[0042] It should be noted that the terms "first", "second", etc. in the description, claims and the above drawings of the present invention are used to distinguish similar objects, and do not necessarily have to be used to describe a specific order or sequence. It should be understood that the data used in this way can be interchanged under appropriate circumstances, so that the embodiments of the present invention described herein can be implemented in an order other than those illustrated or described herein. In addition, the terms "including" and "having" and any variations thereof are intended to cover non-exclusive inclusion. For example, a process, method, system, product or device that includes a series of steps or units does not necessarily have to be limited to those steps or units clearly listed, but may include other steps or units not clearly listed or inherent to these processes, methods, products or devices.

[0043] As Figure 1 shown, the present invention provides an underwater image enhancement method based on a dual-channel convolutional neural network. The dual-channel convolutional neural network includes a shallow convolutional neural network, a deep convolutional neural network, and a fine-tuning convolutional neural network. The underwater image enhancement method includes the following steps:

[0044] Data collection: Collect real underwater image datasets at different sea areas and different depths and synthetic underwater image sets under different scenarios and conditions as the training set and the validation set for this convolutional neural network. Because if only a single underwater image dataset is used as the training set, the generalization ability of the network will be weak. Only by ensuring the diversity of the selected underwater images can the generalization performance of the network be improved;

[0045] Data augmentation: In deep learning, it is generally required that the number of samples be sufficient. The more samples there are, the better the trained model will be. To increase the amount of the selected sample data, the method of data augmentation is used, such as image flipping, image scaling, image cropping, image translation, etc. to double the sample capacity;

[0046] Build a shallow neural network module: The size of the input layer is H*W*3. The first to fifth convolutional layers use a kernel of size 3*3*64, and a residual function is added between the first convolutional layer and the fifth convolutional layer, and finally the first output image of size H*W*64 is obtained;

[0047] Build a deep neural network module: Design a self-attention module (CAB) to extract image feature information. The specific structure is as follows: Two convolutional kernels with a kernel size of 3*3*64 are densely connected through a tensor concatenation (concat) operation in a dense unit, followed by another convolutional kernel of 3*3*64, and the attention mechanism SE Block is used to design weights for each channel. The first convolutional layer of the deep neural network module uses a kernel of size 3*3*64 for preprocessing, and then passes through four CAB modules to obtain the second output image of size H*W*64;

[0048] Build a fine-tuning neural network module: First, use the torch.cat command to concatenate the first output image and the second output image in the channel dimension. The size of the concatenated image is H*W*128. The fine-tuning neural network has six layers. The first convolutional layer uses a kernel of size 3*3*128, the second to fifth layers use kernels of size 3*3*64, and the sixth convolutional layer uses a kernel of 3*3*3 to finally obtain the enhanced underwater image;

[0049] This convolutional neural network uses Leaky ReLU as the activation function, the BatchNormalization function for normalization, the mean squared error (MSE), mean absolute error (MAE), and structural similarity index (SSIM) as the loss functions, and the Adam optimizer.

[0050] The present invention also provides a storage medium, which includes a stored program. When the program runs, it executes the underwater image enhancement method based on the dual-channel convolutional neural network.

[0051] The present invention also provides an electronic device, which includes a memory, a processor, and a computer program stored on the memory and executable on the processor. The processor runs the computer program to execute the underwater image enhancement method based on the dual-channel convolutional neural network.

[0052] The serial numbers of the above embodiments of the present invention are only for description and do not represent the advantages and disadvantages of the embodiments.

[0053] In the above embodiments of the present invention, the descriptions of each embodiment have their own focuses. For the parts not detailed in a certain embodiment, reference can be made to the relevant descriptions of other embodiments.

[0054] In several embodiments provided by this application, it should be understood that the disclosed technical content can be implemented in other ways. Among them, the device embodiments described above are merely illustrative. For example, the division of the units can be a logical function division. In actual implementation, there can be other division methods. For example, multiple units or components can be combined or integrated into another system, or some features can be ignored or not executed. Another point is that the displayed or discussed coupling or direct coupling or communication connection between each other can be through some interfaces. The indirect coupling or communication connection of units or modules can be in an electrical or other form.

[0055] The units described as separate components may or may not be physically separated. The components displayed as units may or may not be physical units, that is, they can be located in one place or distributed to multiple units. Some or all of the units can be selected according to actual needs to achieve the purpose of the solution of this embodiment.

[0056] In addition, in each embodiment of the present invention, the functional units can be integrated in a processing unit, or each unit can exist physically alone, or two or more units can be integrated in one unit. The above-mentioned integrated units can be implemented in the form of hardware or in the form of software functional units.

[0057] If the above-mentioned integrated units are implemented in the form of software functional units and sold or used as independent products, they can be stored in a computer-readable storage medium. Based on this understanding, the technical solution of the present invention, in essence, or the part that contributes to the prior art, or all or part of this technical solution, can be embodied in the form of a software product. This computer software product is stored in a storage medium and includes several instructions for causing a computer device (which can be a personal computer, a server or a network device, etc.) to execute all or part of the steps of the methods described in each embodiment of the present invention. The aforementioned storage medium includes: various media such as USB flash drives, read-only memories (ROM, Read-Only Memory), random access memories (RAM, Random Access Memory), mobile hard disks, magnetic disks or optical discs that can store program codes.

[0058] Finally, it should be noted that: the above embodiments are only used to illustrate the technical solutions of the present invention, rather than to limit it; although the present invention has been described in detail with reference to the foregoing embodiments, those of ordinary skill in the art should understand that: they can still modify the technical solutions recorded in the foregoing embodiments, or perform equivalent replacements on some or all of the technical features; and these modifications or replacements do not make the essence of the corresponding technical solutions deviate from the scope of the technical solutions of each embodiment of the present invention.

Claims

1. An underwater image enhancement method based on a dual-channel convolutional neural network, characterized in that, the dual-channel convolutional neural network includes a shallow convolutional neural network, a deep convolutional neural network, and a fine-tuning convolutional neural network, and the underwater image enhancement method includes the following steps: S1: Establish an image training set, and select an image from the image training set as the target image to be optimized; S2: Input a target image to be optimized into the shallow convolutional neural network branch to optimize the global structure of the target image, and obtain a globally structure-optimized image; the shallow convolutional neural network contains five convolutional layers, the convolutional kernels are all 3, and after each convolutional layer, a normalization processing function and a LeakyReLU activation function are connected, and residuals are added in the first convolutional block and the last convolutional block; S3: Input another target image to be optimized into the deep convolutional neural network branch to recover the detailed information of the target image, and obtain a detailed information-recovered image; the deep convolutional neural network contains thirteen convolutional layers, the convolutional kernels are all 3, and after each convolutional layer, a normalization processing function and a LeakyReLU activation function are connected, and the attention mechanism module SE Block is used to jointly process the target image to be optimized with the convolutional block; S4: Input the globally structure-optimized image and the detailed information-recovered image into the fine-tuning convolutional neural network, perform image fusion using tensor splicing operations, and output the enhanced target image; the fine-tuning convolutional neural network contains six convolutional layers, the convolutional kernels are all 3, and after each convolutional layer, a normalization processing function and a LeakyReLU activation function are connected; S5: Perform the above steps on all images in the image training set until the image enhancement is completed.

2. The underwater image enhancement method based on a dual-channel convolutional neural network according to claim 1, characterized in that, S1 includes the following steps: Select underwater image datasets in different ocean environments and at different depths as the training set, and use data augmentation to increase the data volume; Collect real underwater image datasets in different sea areas and at different depths and synthetic underwater image sets in different scenarios and conditions as the training set and the validation set as the training set and the validation set of the dual-channel convolutional neural network; Adopt a data augmentation method to double the sample capacity of the training set and the validation set.

3. The underwater image enhancement method based on a dual-channel convolutional neural network according to claim 2, characterized in that, the data augmentation method includes image flipping, image scaling, image cropping, and image translation.

4. The underwater image enhancement method based on a dual-channel convolutional neural network according to claim 1, characterized in that, the target image to be optimized is a 3-channel RGB image.

5. The underwater image enhancement method based on a dual-channel convolutional neural network according to claim 1, characterized in that, S2 includes the following steps: Perform preprocessing on the target image to be optimized to reduce the noise impact in the image; Use four convolutional neural network modules to recover the global structure information of the preprocessed target image to be optimized; The first convolutional block of the shallow convolutional neural network is connected to the last convolutional block by using a residual operation, enabling the network to extract deeper features while avoiding gradient vanishing or explosion, and obtaining a globally structurally optimized image.

6. The underwater image enhancement method based on a dual-channel convolutional neural network according to claim 1, wherein, S3 includes the following steps: Perform preprocessing on the target image to be optimized to reduce the influence of irrelevant information in the image; Use four self-attention modules to further extract information from the preprocessed target image to be optimized; Adopt the channel attention mechanism SE Block to fully distinguish effective information from invalid information. SE Block starts from a single image and extracts image features. The feature map dimension of the current feature layer U is [C, H, W], where H represents height, W represents width, and C represents the number of channels; perform average pooling or max pooling on the [H, W] dimension of the feature map, and the size of the pooled feature map is [C, 1, 1], that is, the weight of each channel is obtained; apply the weight to the feature map U[C, H, W], that is, each channel multiplies itself by its respective weight to obtain a detail information restored image.

7. The underwater image enhancement method based on a dual-channel convolutional neural network according to claim 1, wherein, S4 includes the following steps: Fuse the two images obtained after being processed by the shallow and deep convolutional neural network modules together as a new input image; Use six convolutional blocks to perform color correction on the new input image; Output the enhanced underwater image.

8. A storage medium, wherein, the storage medium includes a stored program, wherein when the program runs, it executes the underwater image enhancement method based on a dual-channel convolutional neural network according to any one of claims 1 to 7.

9. An electronic device, including a memory, a processor, and a computer program stored on the memory and executable on the processor, wherein, the processor runs and executes the underwater image enhancement method based on a dual-channel convolutional neural network according to any one of claims 1 to 7 through the computer program.

Images