An underwater real-time image enhancement network system based on multi-scale feature parallel fusion

Abstract

The underwater real-time image enhancement network based on parallel fusion of multi-scale features comprises four multi-scale feature extraction modules, the multi-scale feature extraction module comprises a receptive field enhancement module, a detail optimization module and a CBAM module, the parallel outputs of the receptive field enhancement module and the detail optimization module are added and then processed through the CBAM module to serve as the output of the multi-scale feature extraction module; after the image is subjected to convolution processing, the image is input into the first multi-scale feature extraction module, the output of the current multi-scale feature extraction module is stacked with the original image and the outputs of all the previous stage multi-scale feature extraction modules in the channel dimension to serve as the input of the next stage, and the network model forms a dense connection; after the image is subjected to feature extraction through the multi-scale feature extraction module, the image is output to a main task head and an auxiliary task head, and supervision in different scales is realized. The underwater real-time image enhancement network can well meet the real-time requirement with a small amount of parameters.

Description

Technical Field

[0001] This invention belongs to the field of image processing technology and relates to an underwater real-time image enhancement network system based on multi-scale feature parallel fusion. Background Technology

[0002] Underwater image enhancement utilizes computer vision techniques to improve the quality and visualization of images captured in underwater environments. Due to physical phenomena such as light scattering and refraction in the underwater environment, real underwater images often suffer from low contrast, color distortion, blurriness, and noise. These problems hinder image utilization and further analysis. Underwater image enhancement uses image processing techniques to alleviate these issues, resulting in clearer images with richer information.

[0003] While underwater image enhancement has important applications in various fields, several challenges limit its effectiveness and practical application. Building a physical model and processing images by considering underwater lighting conditions and environmental characteristics is a complex task. Different water areas and depths result in significant differences in image quality, making it difficult to establish a universally applicable physical model. Fluctuations and irregularities in underwater lighting hinder the accurate recovery of image details and texture features. Underwater images often encounter various complex environmental noises, color distortions, and artifacts, further leading to the loss of image information. Overcoming these difficulties and ensuring the acquisition of valuable information is a challenging task. Especially in real-time underwater robot navigation applications, there is an urgent need for fast and effective image enhancement solutions. Solving these demanding problems is a key challenge at present.

[0004] Existing underwater image enhancement methods are mainly divided into three categories: vision-prior-based methods, physics-based methods, and data-driven methods. While these techniques have improved the quality of underwater images to some extent, they still have certain limitations. Vision-prior-based methods use techniques such as color correction, histogram equalization, and wavelet transform to improve image contrast and sharpness. However, such methods still have problems when dealing with complex issues, such as changes in lighting conditions and color distortion. Physics-based techniques mainly model underwater light propagation and water environment factors to reduce the impact of light fluctuations and scattering on images. However, this method requires accurate environmental data and may not be able to fully consider the complex relationships between multiple factors when dealing with complex real-world underwater conditions. Data-driven techniques using convolutional neural networks (CNNs) and generative adversarial networks (GANs) have proven effective in learning underwater image features. They can perform tasks such as denoising and enhancement, and have achieved significant improvements in image quality. While deep learning models used for underwater image enhancement have achieved success, challenges remain. First, the large number of parameters in deep learning models limits their real-time image enhancement applications on embedded devices. Real-time image processing is crucial for underwater robot operations, but the high computational demands hinder the fulfillment of real-time image enhancement requirements. Secondly, complex models require large datasets to fully learn the features of underwater images. However, currently available underwater environment datasets for training are limited, especially high-quality labeled datasets. This data scarcity significantly hinders the training of large-parameter models and impacts their performance and generalization ability. Summary of the Invention

[0005] To address the aforementioned technical problems, the present invention aims to provide an underwater real-time image enhancement network system based on multi-scale feature parallel fusion, which can effectively meet real-time requirements with minimal parameter usage, resulting in more vibrant and saturated colors in the restored images.

[0006] This invention provides an underwater real-time image enhancement network system based on multi-scale feature parallel fusion, comprising four serially connected multi-scale feature extraction modules. Each multi-scale feature extraction module includes a receptive field enhancement module, a detail optimization module, and a CBAM module. The parallel outputs of the receptive field enhancement module and the detail optimization module are summed and processed by the CBAM module before being used as the output of the multi-scale feature extraction module. The image is processed by a 3×3 convolution and then input to the first multi-scale feature extraction module. The output of the current multi-scale feature extraction module is stacked along with the original image and the outputs of all previous multi-scale feature extraction modules in the channel dimension as the input for the next stage, forming a densely connected network model. After feature extraction by the four multi-scale feature extraction modules, the underwater image is output to a main task head and an auxiliary task head, respectively, achieving supervision at different scales.

[0007] Furthermore, the receptive field enhancement module can expand the receptive field of the convolutional kernel, enabling the network to capture more contextual information from the input image;

[0008] After the feature map is input into the receptive field enhancement module, it first passes through a 3×3 convolutional layer with a stride of 2. This reduces the size of the feature map while increasing the number of channels, allowing subsequent convolutional kernels to capture information from a larger spatial range. Then, it is upsampled using simple nearest neighbor interpolation to restore the original size. After reducing the number of channels using a 1×1 convolutional layer, it is stacked with the feature map input to this module in terms of channel dimension. Finally, a 3×3 convolution is used to refine and fuse the features.

[0009] Furthermore, the detail optimization module enables the network model to better acquire local information and enhance the detailed features of the image;

[0010] After the feature map is input into the detail optimization module, it is upsampled to the nearest neighbor, making the height and width of the image twice as big as they were originally. Then, a 1×1 convolutional layer is used to reduce the number of channels, and then a 3×3 convolutional layer with a stride of 2 is used to change the size of the feature map back to its original size. Finally, a 3×3 convolutional layer is used to further refine the tiny features while changing the number of channels.

[0011] Furthermore, the main task head is responsible for training supervision of the same size, while the auxiliary task head performs training supervision on the reduced-size image after shrinking the feature map using a 3×3 convolution kernel with a stride of 2.

[0012] Furthermore, the loss function for training the network model. for:

[0013]

[0014]

[0015] in, and These are hyperparameters, set to 0.7 and 5; This represents the loss of the nth task; where It is the mean squared error loss of the nth task; It is the VGG loss of the nth task; It is the SSIM loss of the nth task; where n=1 represents the main task and n=2 represents the auxiliary task.

[0016] Furthermore, the mean square error loss is calculated using the following formula:

[0017]

[0018] Where B is the batch size, C is the number of channels, and H and W are the height and width of the image, respectively. Generate images for the network. This is a real image.

[0019] Furthermore, the VGG loss is calculated according to the following formula:

[0020]

[0021] Where M represents the number of features extracted from the VGG model. and Let represent the representations of the network-generated image and the real image on the m-th feature, respectively.

[0022] Furthermore, the SSIM loss is calculated according to the following formula:

[0023]

[0024] Where x and y represent two input images, and They are their average values, and Their standard deviations are respectively. Let their covariance be... It is 0.01 2 and It is 0.03 2 They are two constants.

[0025] The underwater real-time image enhancement network system based on multi-scale feature parallel fusion of the present invention has at least the following beneficial effects:

[0026] 1. The network system of this invention extracts features at different image scales and utilizes dense connections to fully leverage the features extracted at different stages of the network model. It employs only a minimum of 0.21MB of parameters to effectively meet the real-time requirements of underwater image enhancement, thus helping to reduce computational costs and improve real-time performance.

[0027] 2. The network system of this invention can achieve an inference speed of up to 52 FPS on the Jetson Orin Nano development board. Its real-time performance and hardware compatibility make it of great value in underwater robot applications.

[0028] 3. The image enhancement network system of this invention can improve the quality of underwater images, which helps to improve the execution efficiency and accuracy of various underwater robot system modules, such as task execution efficiency, navigation accuracy, and decision-making. This correlation makes the image enhancement technology of this invention play a key role in underwater robot systems, improving the overall system performance and making it more suitable for various tasks, including underwater search, exploration, scientific research, and resource management. Attached Figure Description

[0029] Figure 1 This is a structural diagram of an underwater real-time image enhancement network system based on multi-scale feature parallel fusion according to the present invention;

[0030] Figure 2 This is a schematic diagram of the receptive field enhancement module;

[0031] Figure 3 This is a schematic diagram of the detail optimization module. Detailed Implementation

[0032] like Figure 1 As shown, this invention discloses a real-time underwater image enhancement network system based on parallel fusion of multi-scale features. It includes four serially connected multi-scale feature extraction modules: a Receptive Field Enhanced (RFE) module, a Fine-Grained Detail (FGD) module, and a CBAM module. The parallel outputs of the RFE and FGD modules are summed and processed by the CBAM module before being used as the output of the multi-scale feature extraction module. The underwater image, after undergoing 3×3 convolution processing, is input to the first multi-scale feature extraction module. The output of the current multi-scale feature extraction module is stacked along with the original image and the outputs of all previous multi-scale feature extraction modules in the channel dimension as the input for the next stage, creating a densely connected network model. After feature extraction by the four multi-scale feature extraction modules, the underwater image is output to a main task head and an auxiliary task head, respectively, enabling supervision at different scales.

[0033] Before stacking the output of the current multi-scale feature extraction module with the output of the previous multi-scale feature extraction module in the channel dimension, it is necessary to reduce the number of channels of the output of the previous multi-scale feature extraction module by 1×1 convolution before stacking in the channel dimension.

[0034] The densely connected network architecture employed in this invention promotes the flow and reuse of feature information between different stages of the network, enhancing the network model's ability to capture features at different levels. Furthermore, dense connections facilitate gradient backpropagation and, while enhancing image details, effectively preserve the original image's features, resulting in better image enhancement.

[0035] This invention employs a CBAM module to simultaneously apply the attention mechanism in both channel and spatial dimensions. This tells the network model what to pay attention to, while also enhancing the representation of specific regions. It improves the feature extraction capability of the network model without significantly increasing computational cost or the number of parameters.

[0036] like Figure 2 As shown, the receptive field enhancement module of this invention can expand the receptive field of the convolutional kernel, enabling the network to capture more contextual information from the input image. After the feature map is input into the receptive field enhancement module, it first passes through a 3×3 convolutional layer with a stride of 2, which reduces the feature map size while increasing the number of channels, allowing subsequent convolutional kernels to acquire information over a larger spatial range. Then, simple nearest neighbor interpolation is used for upsampling to restore the original size. A 1×1 convolutional layer is then used to reduce the number of channels and stacked with the feature map input to this module in the channel dimension, achieving a residual connection effect. Finally, a 3×3 convolution is used to further refine and fuse the features.

[0037] like Figure 3 As shown, the detail optimization module enables the network model to better acquire local information, thereby enhancing the detailed features of the image. After the feature map is input into the detail optimization module, nearest neighbor upsampling is performed, doubling both the height and width of the image. Then, a 1×1 convolutional layer reduces the number of channels, minimizing subsequent computational overhead. Next, a 3×3 convolutional layer with a stride of 2 resizes the feature map back to its original size. Finally, a 3×3 convolutional layer further refines minute features while changing the number of channels. The detail optimization module of this invention not only improves the network model's ability to extract local features but also maintains high computational efficiency throughout the process.

[0038] In this invention, a multi-task head is employed for image enhancement tasks. One main task head handles training supervision of images of the same size, facilitating the reconstruction of image details. Another auxiliary task head, after shrinking the feature map using a 3×3 convolutional kernel with a stride of 2, supervises images of size H / 2 and W / 2, where H and W are the height and width of the feature map. By using a multi-task head, the receptive field of the convolutional kernel is enhanced, allowing for the acquisition of more regional information for feedback, which is beneficial for better model optimization during gradient backpropagation.

[0039] To better enhance image details, addressing the challenge of image detail enhancement, MSE loss is used for pixel-level evaluation, and L1 loss is used to evaluate the semantic information extracted by the VGG model. The VGG model is a classic deep learning model frequently used to extract image features. Furthermore, considering the differences in brightness, contrast, and other information across different image regions, a block-based SSIM loss function is employed. These three functions together constitute the loss function for a single task. The overall mathematical expression for the loss function during network model training is:

[0040]

[0041]

[0042] in, This is the loss function of the network model in this invention. and These are hyperparameters, set to 0.7 and 5; This represents the loss of the nth task; where It is the mean squared error loss of the nth task; It is the VGG loss of the nth task; It is the SSIM loss of the nth task; where n=1 represents the main task and n=2 represents the auxiliary task.

[0043] In practical implementation, the mean square error loss is calculated according to the following formula:

[0044]

[0045] Where B is the batch size, C is the number of channels, and H and W are the height and width of the image, respectively. Generate images for the network. This is a real image.

[0046] In practical implementation, VGG loss is calculated according to the following formula:

[0047]

[0048] Where M represents the number of features extracted from the VGG model. and Let represent the representations of the network-generated image and the real image on the m-th feature, respectively.

[0049] In practice, the SSIM loss is calculated using the following formula:

[0050]

[0051] Where x and y represent two input images, and They are their average values, and Their standard deviations are respectively. Let their covariance be... It is 0.01 2 and It is 0.03 2 They are two constants.

[0052] The network system of this invention was compared with WaterNet, U-shaped Transformer, Shallow-Uwnet, and DeepWaveNet models. All models in the comparison experiments used the adam optimizer. Considering the limited computing power of edge computing devices and the high real-time requirements for image processing in real-world environments, the batch size was set to 1, the maximum learning rate was 0.0002, and the system was trained for 100 epochs. The first five epochs used a warmup to gradually increase the learning rate, followed by cosine annealing to gradually decrease the learning rate to 1e-5. The model was trained using the PyTorch framework on a server with an RTX 3090 and an Intel(R) Xeon(R) CPU E5-2678v3. Comparative experiments were conducted on the three EUVP sub-data sets to demonstrate the superiority of the network model of this invention. Specifically, considering that the training parameters of the U-shaped Transformer are much larger than those of other models, the officially recommended maximum learning rate of 0.0005 was used for the U-shaped Transformer, while other training parameters remained unchanged. The results of the comparative experiments are shown in Table 1.

[0053] Table 1 Comparative Experiments

[0054]

[0055] As can be seen from Table 1, the network system of this invention achieves the best PSNR and SSIM scores on the three EUVP subsets simultaneously, and significantly outperforms other models, demonstrating its superiority and strong fitting ability to different feature data.

[0056] A significant feature of the network system of this invention is its ability to be deployed on embedded devices. The network model employs a convolutional neural network with only 0.21 MB of parameters, exhibiting low computational complexity and memory footprint. Further experiments were conducted using TensorRT with FP16 precision on a Jetson Orin Nano development board, and FPS test results for different versions of the model under a batch size of 1 are provided. The model's input and output sizes are both 256×256. Table 2 shows the model deployment test results.

[0057] Table 2 Model Deployment Testing

[0058]

[0059] Table 2 shows the performance of different versions of the network model. The more complex the modules used in the model, the more parameters and computational cost increase, resulting in a slight decrease in frame rate while improving performance. Experimental results show that the network model of this invention can achieve a maximum inference capability of 52 FPS, and its accuracy surpasses existing advanced models. It can fully meet the real-time detection requirements of edge computing devices at a high accuracy level, reducing the deployment cost of the model. Other models are designed for scenarios with high image quality requirements and can also meet the real-time processing needs of most application scenarios.

[0060] Current traditional physical models and vision-based prior methods for underwater image enhancement cannot adapt well to different underwater environments. Their effectiveness is greatly affected by ambient lighting and noise, and they are highly dependent on hyperparameters. Real-world applications place high demands on the generalization ability of these models. The underwater real-time image enhancement network of this invention can learn underwater environment-specific features and patterns from a large amount of underwater image data, thus better adapting to problems such as light attenuation and scattering, and automatically adjusting image enhancement parameters to achieve more accurate image restoration. Tasks such as underwater robot operation and underwater resource exploration place extremely high demands on rapid feedback and decision-making. However, current deep learning models have excessively large parameter counts and computational costs, hindering their deployment on embedded devices. The underwater real-time enhancement network of this invention can effectively meet real-time requirements with a minimal number of parameters.

[0061] The above description is only a preferred embodiment of the present invention and is not intended to limit the ideas of the present invention. Any modifications, equivalent substitutions, improvements, etc., made within the spirit and principles of the present invention should be included within the protection scope of the present invention.

Claims

1. A real-time underwater image enhancement network system based on multi-scale feature parallel fusion, characterized in that, It includes four serially connected multi-scale feature extraction modules, each comprising a receptive field enhancement module, a detail optimization module, and a CBAM module. The parallel outputs of the receptive field enhancement module and the detail optimization module are added together and then processed by the CBAM module to serve as the output of the multi-scale feature extraction module. After the image is processed by 3×3 convolution, it is input into the first multi-scale feature extraction module. The output of the current multi-scale feature extraction module is stacked with the original image and the outputs of all previous multi-scale feature extraction modules in the channel dimension as the input of the next stage, so that the entire network model forms a dense connection. After the underwater image is processed by the feature extraction of four multi-scale feature extraction modules, it is output to a main task head and an auxiliary task head respectively, realizing supervision at different scales. The receptive field enhancement module can expand the receptive field of the convolution kernel, enabling the network to capture more contextual information from the input image; After the feature map is input into the receptive field enhancement module, it first passes through a 3×3 convolutional layer with a stride of 2. This reduces the feature map size while increasing the number of channels, allowing subsequent convolutional kernels to capture information from a larger spatial range. Then, it is upsampled using simple nearest neighbor interpolation to restore the original size. After reducing the number of channels using a 1×1 convolutional layer, it is stacked with the feature map input to this module in terms of channel dimension. Finally, a 3×3 convolution is used to refine and fuse the features. The detail optimization module enables the network model to better acquire local information and enhances the detailed features of the image. After the feature map is input into the detail optimization module, it is upsampled to the nearest neighbor, making the height and width of the image twice as big as they were originally. Then, a 1×1 convolutional layer is used to reduce the number of channels, and then a 3×3 convolutional layer with a stride of 2 is used to change the size of the feature map back to its original size. Finally, a 3×3 convolutional layer is used to further refine the tiny features while changing the number of channels.

2. The underwater real-time image enhancement network system based on multi-scale feature parallel fusion as described in claim 1, characterized in that, The main task head is responsible for training supervision of images of the same size, while the auxiliary task head performs training supervision on images of the reduced size after shrinking the feature map using a 3×3 convolutional kernel with a stride of 2.

3. The underwater real-time image enhancement network system based on multi-scale feature parallel fusion as described in claim 1, characterized in that, Loss function for training network models for: in, and These are hyperparameters, set to 0.7 and 5; This represents the loss of the nth task; where It is the mean squared error loss of the nth task; It is the VGG loss of the nth task; It is the SSIM loss of the nth task; where n=1 represents the main task and n=2 represents the auxiliary task.

4. The underwater real-time image enhancement network system based on multi-scale feature parallel fusion as described in claim 3, characterized in that, The mean square error loss is calculated using the following formula: Where B is the batch size, C is the number of channels, and H and W are the height and width of the image, respectively. Generate images for the network. This is a real image.

5. The underwater real-time image enhancement network system based on multi-scale feature parallel fusion as described in claim 4, characterized in that, VGG loss is calculated using the following formula: Where M represents the number of features extracted from the VGG model. and Let represent the representations of the network-generated image and the real image on the m-th feature, respectively.

6. The underwater real-time image enhancement network system based on multi-scale feature parallel fusion as described in claim 3, characterized in that, The SSIM loss is calculated using the following formula: Where x and y represent two input images, and They are their average values, and Their standard deviations are respectively. Let their covariance be... It is 0.01 2 and It is 0.03 2 They are two constants.

Images