Underwater image enhancement method and device based on convolutional neural network

Abstract

The invention discloses an underwater image enhancement method and device based on a convolutional neural network, and belongs to the technical field of machine learning and image vision. The method comprises the following steps: S1, data preprocessing, including image data collection, image normalization, image color correction and image cutting; s2, designing a network architecture, and adding a self-attention mechanism in the network architecture; s3, designing a comprehensive loss function; s4, performing model training, dividing the acquired image data into a training set, a verification set and a test set, and performing training by using the comprehensive loss function; s5, verifying and testing the trained model; and S6, optimizing and deploying the model, adopting a model compression, pruning and quantification technology for an optimization strategy, and starting to deploy the model for practical application after the optimization is completed. According to the method, a comprehensive loss function is designed, perception loss and antagonistic loss are combined, and an image enhancement effect is achieved through comprehensive evaluation and optimization.

Description

Technical Field

[0001] The present invention relates to the field of machine learning and image vision technology, and in particular to an underwater image enhancement method and device based on a convolutional neural network. Background Art

[0002] Convolutional Neural Network (CNN) is a deep learning model widely used in image and video recognition. CNN automatically extracts and learns the features of input data by using convolutional layers, pooling layers, and fully connected layers.

[0003] At present, underwater images are widely used, especially in underwater biological identification, underwater environmental exploration, underwater cultural relics excavation and underwater environmental development. However, due to the lack of underwater light, turbid water and other objective conditions, underwater images are blurred, color distorted, texture blurred and lack contrast, which cannot provide effective help for current underwater exploration. At present, machine learning is widely used in processing a large number of underwater images, but the current machine learning models are conventional methods, which are self-learned through a large amount of image data. This self-learning method takes a lot of resources to capture feature information of different scales, and it is difficult to capture long-distance dependent information, which leads to problems in feature extraction of the model and cannot be accurately extracted; the current model needs to be trained through a large amount of image data annotation to improve the generalization ability of the model, which wastes a lot of manpower, training time and resources, and the model does not consider the loss calculation comprehensively, resulting in differences in model optimization and unstable image enhancement effect. Summary of the invention

[0004] The purpose of the present invention is to provide an underwater image enhancement method and device based on convolutional neural network to solve the problems raised in the above background technology.

[0005] In order to solve the above technical problems, the present invention provides the following technical solutions: an underwater image enhancement method based on a convolutional neural network, the method comprising the following steps:

[0006] S1, data preprocessing, including image data collection, image normalization, image color correction and image cropping;

[0007] S2. Network architecture design. ResNet is selected as the basic network architecture for initial feature extraction and image reconstruction. A self-attention mechanism is added to the network, and a feature pyramid network FPN structure is used to fuse features of different scales.

[0008] S3. Loss function design: combining perceptual loss and adversarial loss to design a comprehensive loss function;

[0009] S4, model training, dividing the collected image data into a training set, a validation set and a test set, and training using the above-mentioned comprehensive loss function;

[0010] S5. Use the validation data set in the image data collection to perform a validation test on the trained model and calculate the PSNR objective evaluation index. If the index is not met, use the validation data set for secondary training until the PSNR objective evaluation index is met.

[0011] S6. Model optimization and deployment. The optimization strategy uses model compression, pruning and quantization techniques. After the optimization is completed, the model is deployed for practical application.

[0012] The image data collection includes actual captured images, AI synthesized images, and public underwater image datasets on github;

[0013] The actual captured images include images of multiple different water areas with a depth range of [1,70] meters; the AI ​​synthetic images are mainly generated images with specific features, including contrast, color difference, lighting and texture;

[0014] The image dataset consists of the above three parts of images and is used for model training, verification and testing.

[0015] In step S1,

[0016] Image normalization is to normalize the pixel values ​​of the image to the range [-1, 1], which is calculated by subtracting the image pixel mean from the pixel value of the image and dividing it by the standard deviation;

[0017] Image color correction is used to correct color casts by adjusting the gain of each color channel so that the visual quality of the target image is consistent with the reference image input to the network.

[0018] The image cropping is used to crop and scale the image to a uniform size to ensure that the size of the input image is consistent; the image is scaled to a uniform size, and the specific image size is 256×256.

[0019] In step S2, the basic network architecture ResNet includes an input layer, a convolutional layer, a residual module, and a fully connected layer;

[0020] The input layer is used to input an image, that is, the image of size 256×256 after the preprocessing; the convolution layer includes multiple convolution layers Conv and batch normalization BatchNorm layers, which are used to extract low-level features; the residual module includes multiple residual blocks, each residual block includes two convolution layers and a batch normalization layer; the fully connected layer is used to connect the above residual modules and output the final feature map.

[0021] The self-attention mechanism is inserted into the middle layer of ResNet. The specific operation of the self-attention module is as follows:

[0022] S5-1, input feature map from the middle layer of ResNet, size is H×W×C;

[0023] S5-2, perform transformation operations, generate a Query feature map through 1×1 convolution, generate a Key feature map through 1×1 convolution, and generate a Value feature map through 1×1 convolution;

[0024] S5-3, similarity calculation, calculate the dot product similarity between Query and Key to obtain the attention weight matrix;

[0025] S5-4, weighted summation, multiply the attention weight matrix by the Value feature map to obtain a weighted sum feature map;

[0026] S5-5, output feature map, add the weighted sum feature map to the original input feature map to obtain an enhanced feature map.

[0027] The feature pyramid network FPN is used for multi-scale feature fusion, and performs feature fusion in two ways: bottom-up and top-down. Specifically:

[0028] The bottom-up path of the FPN structure outputs feature maps of different scales from each stage of ResNet;

[0029] The top-down path of the FPN structure starts from the highest-level feature map, gradually upsamples downward, and merges with the bottom-up feature map. The number of high-level feature map channels is halved through 1×1 convolution, and then added to the corresponding low-level feature map through 2×2 upsampling.

[0030] The feature fusion is to further process the fused feature map through 3×3 convolution to obtain the final multi-scale feature map.

[0031] In step S3, the comprehensive loss function includes mean square error loss, structural similarity loss, perceptual loss and adversarial loss;

[0032] The mean square error loss:

[0033]

[0034] Among them, MSE Lo ss represents the mean square error loss; yi represents the i-th pixel value of the real image; Represents the i-th pixel value of the predicted image, 1≤i≤n; n represents the total number of pixels in the image;

[0035] The structural similarity loss:

[0036]

[0037] in, represents the structural similarity loss; μ y and Represent the pixel averages of the real image and the predicted image respectively; σ y and Represent the variance of the pixel values ​​of the real image and the predicted image respectively; represents the covariance of the real image and the predicted image; C1 and C2 are structural constants;

[0038] The perceived loss:

[0039]

[0040] Among them, Perceptual Loss Indicates perceived loss; express The L2 norm of , which represents the Euclidean distance; represents the j-th layer mapping of the training network; y and Represent the real image and the predicted image respectively;

[0041] The adversarial loss:

[0042] Adversarial Loss =E x～Pdata [logD(x)]+E z～Pz [log(1-D(G(z)))]

[0043] Among them, Adversarial Loss represents the image adversarial loss; E x～Pdata represents the expectation of the real data x sampled from the real data distribution Pdata; E: ~ P z represents the expectation of random noise z sampled from the noise distribution Pz; D and G represent the discriminator and generator in the training network respectively; logD(x) is the value of the discriminator's output for the real data x after being processed by the logarithmic function; this term is expected to be maximized so that the discriminator can correctly classify real images; log(1-D(G(z))) is the value of the discriminator's output for the generated data G(z) after being processed by the logarithmic function; this term is expected to be minimized so that the fake images generated by the generator can deceive the discriminator and make the discriminator think they are real.

[0044] The comprehensive loss function:

[0045]

[0046] Among them, Total Loss represents the comprehensive loss function; α, β, γ and δ are the weight parameters of mean square error loss, structural similarity loss, perceptual loss and adversarial loss respectively;

[0047] In step S5,

[0048] The PSNR objective evaluation index is calculated as follows:

[0049] PSNR is based on mean square error (MSE) Loss The specific formula for calculation is:

[0050]

[0051] Among them, MSE Loss Represents mean square error; PSNR is an objective evaluation indicator; MAX I is the maximum pixel value of the image;

[0052] The verification test is to test the trained model using the verification data set, generate enhanced images, and then calculate the PSNR objective evaluation index for each pair of original images and enhanced images; calculate the average PSNR value on the verification set; if it is lower than the set threshold, perform secondary training;

[0053] Secondary training is to add the validation data set to the training data and retrain with different hyperparameters; retrain the model until the PSNR reaches the preset indicator;

[0054] The threshold is determined by calculating the mean value, median and standard deviation of the error between the reference target image and the real image.

[0055] An underwater image enhancement device based on a convolutional neural network, the device comprising a data preprocessing module, a network architecture design module, a loss function design module, a model training module, a verification test module and a model optimization and deployment module;

[0056] The data preprocessing module is used for image data collection, image normalization, image color correction and image cropping; collects a variety of underwater image data, including actual captured images, AI synthesized images and public underwater image data sets; normalizes the pixel values ​​of the image to the range of [-1,1] and adjusts the gain of each color channel; crops and scales the image to a uniform size to ensure that the size of the input image is consistent;

[0057] The network architecture design module selects the ResNet architecture as the basic network architecture for initial feature extraction and image reconstruction, introduces a self-attention mechanism module in the middle layer of ResNet, and uses the feature pyramid network FPN for multi-scale feature fusion;

[0058] The loss function design module is used to design a comprehensive loss function, which combines mean square error loss, structural similarity loss, perceptual loss and adversarial loss;

[0059] The model training module is used to divide the image data into a training set, a validation set and a test set, and train the convolutional neural network model using a designed comprehensive loss function;

[0060] The verification test module is used to verify the model performance, use the verification data set to test the trained model, calculate the average PSNR value on the verification set, and determine whether the model meets the preset performance standard. If the PSNR does not reach the set threshold, use the verification data set for secondary training;

[0061] The model optimization and deployment module is used to optimize and deploy the trained model. Model optimization includes model compression, model pruning and model quantization. Model deployment is to deploy the model to actual applications after model optimization is completed.

[0062] The output end of the data preprocessing module is connected to the input end of the network architecture design module; the output end of the network architecture design module is connected to the input end of the loss function design module; the output end of the loss function design module is connected to the input end of the model training module; the output end of the model training module is connected to the input end of the verification and testing module; the output end of the verification and testing module is connected to the input end of the model optimization and deployment module.

[0063] Compared with the prior art, the beneficial effects achieved by the present invention are:

[0064] 1. The present invention optimizes the existing network architecture, adds a self-attention mechanism on the basis of the existing network structure, and adopts multi-scale feature fusion to capture feature information of different scales, including long-distance dependencies and global information, thereby improving the feature extraction and enhancement capabilities of the model;

[0065] 2. The present invention adds a self-supervised learning method to reduce the reliance on a large amount of labeled data and improve the generalization ability of the model;

[0066] 3. The present invention designs a comprehensive loss function that combines perceptual loss and adversarial loss to achieve image enhancement effect through comprehensive evaluation and optimization. BRIEF DESCRIPTION OF THE DRAWINGS

[0067] The accompanying drawings are used to provide a further understanding of the present invention and constitute a part of the specification. Together with the embodiments of the present invention, they are used to explain the present invention and do not constitute a limitation of the present invention. In the accompanying drawings:

[0068] Figure 1 It is a schematic diagram of the steps of an underwater image enhancement method based on a convolutional neural network of the present invention;

[0069] Figure 2 It is a schematic diagram of the network architecture of the present invention. DETAILED DESCRIPTION

[0070] The following will be combined with the drawings in the embodiments of the present invention to clearly and completely describe the technical solutions in the embodiments of the present invention. Obviously, the described embodiments are only part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without creative work are within the scope of protection of the present invention.

[0071] See also Figure 1-Figure 2 The present invention provides a technical solution: an underwater image enhancement method based on a convolutional neural network, the method comprising the following steps:

[0072] S1, data preprocessing, including image data collection, image normalization, image color correction and image cropping;

[0073] S2. Network architecture design. ResNet is selected as the basic network architecture for initial feature extraction and image reconstruction. A self-attention mechanism is added to the network, and a feature pyramid network FPN structure is used to fuse features of different scales.

[0074] S3. Loss function design: combining perceptual loss and adversarial loss to design a comprehensive loss function;

[0075] S4, model training, dividing the collected image data into a training set, a validation set and a test set, and training using the above-mentioned comprehensive loss function;

[0076] S5. Use the validation data set in the image data collection to perform a validation test on the trained model and calculate the PSNR objective evaluation index. If the index is not met, use the validation data set for secondary training until the PSNR objective evaluation index is met.

[0077] S6. Model optimization and deployment. The optimization strategy uses model compression, pruning and quantization techniques. After the optimization is completed, the model is deployed for practical application.

[0078] The image data collection includes actual captured images, AI synthesized images, and public underwater image datasets on github;

[0079] The actual captured images include images of multiple different water areas with a depth range of [1,70] meters; the AI ​​synthetic images are mainly generated images with specific features, including contrast, color difference, lighting and texture;

[0080] The image dataset consists of the above three parts of images and is used for model training, verification and testing.

[0081] In step S1,

[0082] Image normalization is to normalize the pixel values ​​of the image to the range [-1, 1], which is calculated by subtracting the image pixel mean from the pixel value of the image and dividing it by the standard deviation;

[0083] Image color correction is used to correct color casts by adjusting the gain of each color channel so that the visual quality of the target image is consistent with the reference image input to the network.

[0084] The image cropping is used to crop and scale the image to a uniform size to ensure that the size of the input image is consistent; the image is scaled to a uniform size, and the specific image size is 256×256.

[0085] In step S2, the basic network architecture ResNet includes an input layer, a convolutional layer, a residual module, and a fully connected layer;

[0086] The input layer is used to input an image, that is, the image of size 256×256 after the preprocessing; the convolution layer includes multiple convolution layers Conv and batch normalization BatchNorm layers, which are used to extract low-level features; the residual module includes multiple residual blocks, each residual block includes two convolution layers and a batch normalization layer; the fully connected layer is used to connect the above residual modules and output the final feature map.

[0087] The self-attention mechanism is inserted into the middle layer of ResNet. The specific operation of the self-attention module is as follows:

[0088] S5-1, input feature map from the middle layer of ResNet, size is H×W×C;

[0089] S5-2, perform transformation operations, generate a Query feature map through 1×1 convolution, generate a Key feature map through 1×1 convolution, and generate a Value feature map through 1×1 convolution;

[0090] S5-3, similarity calculation, calculate the dot product similarity between Query and Key to obtain the attention weight matrix;

[0091] S5-4, weighted summation, multiply the attention weight matrix by the Value feature map to obtain a weighted sum feature map;

[0092] S5-5, output feature map, add the weighted sum feature map to the original input feature map to obtain an enhanced feature map.

[0093] The feature pyramid network FPN is used for multi-scale feature fusion, and performs feature fusion in two ways: bottom-up and top-down. Specifically:

[0094] The bottom-up path of the FPN structure outputs feature maps of different scales from each stage of ResNet;

[0095] The top-down path of the FPN structure starts from the highest-level feature map, gradually upsamples downward, and merges with the bottom-up feature map. The number of high-level feature map channels is halved through 1×1 convolution, and then added to the corresponding low-level feature map through 2×2 upsampling.

[0096] The feature fusion is to further process the fused feature map through 3×3 convolution to obtain the final multi-scale feature map.

[0097] In step S3, the comprehensive loss function includes mean square error loss, structural similarity loss, perceptual loss and adversarial loss;

[0098] The mean square error loss:

[0099]

[0100] Among them, MSE Lo ss represents the mean square error loss; yi represents the i-th pixel value of the real image; Represents the i-th pixel value of the predicted image, 1≤i≤n; n represents the total number of pixels in the image;

[0101] The structural similarity loss:

[0102]

[0103] in, represents the structural similarity loss; μ y and Represent the pixel averages of the real image and the predicted image respectively; σ y and Represent the variance of the pixel values ​​of the real image and the predicted image respectively; represents the covariance of the real image and the predicted image; C1 and C2 are structural constants;

[0104] The perceived loss:

[0105]

[0106] Among them, Perceptual Loss Indicates perceived loss; express The L2 norm of , which represents the Euclidean distance; represents the j-th layer mapping of the training network; y and Represent the real image and the predicted image respectively;

[0107] The adversarial loss:

[0108] Adversarial Loss =E x～Pdata [logD(x)]+E z～Pz [log(1-D(G(z)))]

[0109] Among them, Adversarial Loss represents the image adversarial loss; E x～Pdata represents the expectation of the real data x sampled from the real data distribution Pdata; E: ~ P z represents the expectation of random noise z sampled from the noise distribution Pz; D and G represent the discriminator and generator in the training network respectively; logD(x) is the value of the discriminator's output for the real data x after being processed by the logarithmic function; this term is expected to be maximized so that the discriminator can correctly classify real images; log(1-D(G(z))) is the value of the discriminator's output for the generated data G(z) after being processed by the logarithmic function; this term is expected to be minimized so that the fake images generated by the generator can deceive the discriminator and make the discriminator think they are real.

[0110] The comprehensive loss function:

[0111]

[0112] Among them, Total Loss represents the comprehensive loss function; α, β, γ and δ are the weight parameters of mean square error loss, structural similarity loss, perceptual loss and adversarial loss respectively;

[0113] In step S5,

[0114] The PSNR objective evaluation index is calculated as follows:

[0115] PSNR is based on mean square error (MSE) Loss The specific formula for calculation is:

[0116]

[0117] Among them, MSE Loss Represents mean square error; PSNR is an objective evaluation indicator; MAX I is the maximum pixel value of the image;

[0118] The verification test is to test the trained model using the verification data set, generate enhanced images, and then calculate the PSNR objective evaluation index for each pair of original images and enhanced images; calculate the average PSNR value on the verification set; if it is lower than the set threshold, perform secondary training;

[0119] Secondary training is to add the validation data set to the training data and retrain with different hyperparameters; retrain the model until the PSNR reaches the preset indicator;

[0120] The threshold is determined by calculating the mean value, median and standard deviation of the error between the reference target image and the real image.

[0121] An underwater image enhancement device based on a convolutional neural network, the device comprising a data preprocessing module, a network architecture design module, a loss function design module, a model training module, a verification test module and a model optimization and deployment module;

[0122] The data preprocessing module is used for image data collection, image normalization, image color correction and image cropping; collects a variety of underwater image data, including actual captured images, AI synthesized images and public underwater image data sets; normalizes the pixel values ​​of the image to the range of [-1,1] and adjusts the gain of each color channel; crops and scales the image to a uniform size to ensure that the size of the input image is consistent;

[0123] The network architecture design module selects the ResNet architecture as the basic network architecture for initial feature extraction and image reconstruction, introduces a self-attention mechanism module in the middle layer of ResNet, and uses the feature pyramid network FPN for multi-scale feature fusion;

[0124] The loss function design module is used to design a comprehensive loss function, which combines mean square error loss, structural similarity loss, perceptual loss and adversarial loss;

[0125] The model training module is used to divide the image data into a training set, a validation set and a test set, and train the convolutional neural network model using a designed comprehensive loss function;

[0126] The verification test module is used to verify the model performance, use the verification data set to test the trained model, calculate the average PSNR value on the verification set, and determine whether the model meets the preset performance standard. If the PSNR does not reach the set threshold, use the verification data set for secondary training;

[0127] The model optimization and deployment module is used to optimize and deploy the trained model. Model optimization includes model compression, model pruning and model quantization. Model deployment is to deploy the model to actual applications after model optimization is completed.

[0128] The output end of the data preprocessing module is connected to the input end of the network architecture design module; the output end of the network architecture design module is connected to the input end of the loss function design module; the output end of the loss function design module is connected to the input end of the model training module; the output end of the model training module is connected to the input end of the verification and testing module; the output end of the verification and testing module is connected to the input end of the model optimization and deployment module.

[0129] In an embodiment of the present invention,

[0130] Data preprocessing module:

[0131] Data Collection:

[0132] Collect 3,000 images of different underwater environments, lighting conditions, and depths, including actual captured images, AI synthesized images, and public underwater image datasets.

[0133] Image normalization:

[0134] Normalize the pixel values ​​of the image to the range [-1,1].

[0135]

[0136] Color correction by adjusting the gain of each color channel;

[0137] The image is (R, G, B) = (150, 120, 100), and after correction it is:

[0138] (R′, G′, B′)=(160,125,110)

[0139] Image cropping and scaling:

[0140] All images were cropped and scaled to 256x256 size.

[0141] Network architecture design module:

[0142] ResNet architecture: ResNet-18 was selected as the basic network architecture for initial feature extraction and image reconstruction.

[0143] Self-attention mechanism module: The self-attention mechanism is introduced in the third layer of ResNet to enhance feature representation.

[0144] Feature pyramid network FPN structure:

[0145] Feature maps of different scales are output from each stage of ResNet, and multi-scale feature fusion is performed through bottom-up and top-down paths.

[0146] Loss function design module:

[0147] Comprehensive loss function:

[0148] Processing 5 pixels, the pixel values ​​of the real image and the predicted image are [100, 150, 200, 250, 300] and [105, 148, 198, 252, 295] respectively, and the calculation is as follows:

[0149] The mean square error loss:

[0150]

[0151] Among them, MSE L oss represents the mean square error loss; yi represents the i-th pixel value of the real image; Represents the i-th pixel value of the predicted image, 1≤i≤n; n represents the total number of pixels in the image;

[0152]

[0153] The structural similarity loss:

[0154]

[0155] in, represents the structural similarity loss; μ y and Represent the pixel averages of the real image and the predicted image respectively; σ y and Represent the variance of the pixel values ​​of the real image and the predicted image respectively; represents the covariance of the real image and the predicted image; C1 and C2 are structural constants; C1 = 0.01; C2 = 0.03;

[0156] μ y =200, σ y =50,

[0157]

[0158] The perceived loss:

[0159]

[0160] Among them, Perceptual Loss Indicates perceived loss; express The L2 norm of , which represents the Euclidean distance; represents the j-th layer mapping of the training network; y and Represent the real image and the predicted image respectively;

[0161] Mapping features of layer j Calculate the perceptual loss:

[0162]

[0163] The adversarial loss:

[0164] Adversarial Loss =E x～Pdata [logD(x)]+E z～Pz [log(1-D(G(z)))]

[0165] Among them, Adversarial Loss represents the image adversarial loss; E x～Pdata represents the expectation of the real data x sampled from the real data distribution Pdata; E z～Pz represents the expectation of random noise z sampled from the noise distribution Pz; D and G represent the discriminator and generator in the training network respectively; logD(x) is the value of the discriminator's output on the real data x after being processed by the logarithmic function; log(1-D(G(z))) is the value of the discriminator's output on the generated data G(z) after being processed by the logarithmic function;

[0166] The discriminator output of the real data is 0.9, and the discriminator output of the generated data is 0.2. The adversarial loss is calculated as follows:

[0167] Adversarial Loss =log(0.9)+log(1-0.2)≈-0.3285

[0168] The absolute value of the error is Adcersarial Loss =0.3258

[0169] The comprehensive loss function:

[0170]

[0171] Among them, Total Loss represents the comprehensive loss function; α, β, γ and 6 are the weight parameters of mean square error loss, structural similarity loss, perceptual loss and adversarial loss respectively;

[0172] According to multiple validations, the weights are set to α = 0.2, β = 0.3, γ = 0.1 and 6 = 0.4;

[0173] Total Loss=0.2*12.4+0.3*0.203+0.1*0.2+0.4*0.3285

[0174] Model training module:

[0175] Dataset division: training set: 2400 images, validation set: 400 images, test set: 200 images

[0176] Model training: Use a comprehensive loss function to train the model and optimize model parameters.

[0177] Verification test module:

[0178] Validate model performance: Generate augmented images using the validation dataset.

[0179] Calculate the PSNR value:

[0180]

[0181] The calculated threshold is 40dB. Since it meets the requirements, no secondary training is performed;

[0182] Model optimization and deployment module:

[0183] Model compression: Use knowledge distillation to transfer the knowledge of the original model into a smaller model.

[0184] Model pruning: Pruning neuron connections whose weights are less than a certain threshold.

[0185] Model quantization: Quantize model parameters from 32-bit floating point numbers to 8-bit integers to reduce storage and computing costs.

[0186] Model deployment: Deploy the optimized model to practical applications to improve the quality of underwater images.

[0187] It should be noted that, in this article, relational terms such as first and second, etc. are only used to distinguish one entity or operation from another entity or operation, and do not necessarily require or imply any such actual relationship or order between these entities or operations. Moreover, the terms "include", "comprise" or any other variants thereof are intended to cover non-exclusive inclusion, so that a process, method, article or device including a series of elements includes not only those elements, but also other elements not explicitly listed, or also includes elements inherent to such process, method, article or device.

[0188] Finally, it should be noted that the above description is only a preferred embodiment of the present invention and is not intended to limit the present invention. Although the present invention has been described in detail with reference to the aforementioned embodiments, those skilled in the art can still modify the technical solutions described in the aforementioned embodiments or replace some of the technical features therein by equivalents. Any modification, equivalent replacement, improvement, etc. made within the spirit and principle of the present invention shall be included in the protection scope of the present invention.

Claims

1. An underwater image enhancement method based on convolutional neural network, characterized in that: The method comprises the following steps: S1, data preprocessing, including image data collection, image normalization, image color correction and image cropping; S2. Network architecture design. ResNet is selected as the basic network architecture for initial feature extraction and image reconstruction. A self-attention mechanism is added to the network, and a feature pyramid network FPN structure is used to fuse features of different scales. S3. Loss function design: combining perceptual loss and adversarial loss to design a comprehensive loss function; S4, model training, dividing the collected image data into a training set, a validation set and a test set, and training using the above-mentioned comprehensive loss function; S5. Use the validation data set in the image data collection to perform a validation test on the trained model and calculate the PSNR objective evaluation index. If the index is not met, use the validation data set for secondary training until the PSNR objective evaluation index is met. S6. Model optimization and deployment. The optimization strategy uses model compression, pruning and quantization techniques. After the optimization is completed, the model is deployed for practical application.

2. The underwater image enhancement method based on convolutional neural network according to claim 1, characterized in that: The image data collection includes actual captured images, AI synthesized images, and public underwater image datasets on github; The actual captured images include images of multiple different water areas with a depth range of [1,70] meters; the AI ​​synthetic images are mainly generated images with specific features, including contrast, color difference, lighting and texture; The image dataset consists of the above three parts of images and is used for model training, verification and testing.

3. The underwater image enhancement method based on convolutional neural network according to claim 1, characterized in that: In step S1, Image normalization is to normalize the pixel values ​​of the image to the range [-1, 1], which is calculated by subtracting the image pixel mean from the pixel value of the image and dividing it by the standard deviation; Image color correction is used to correct color casts by adjusting the gain of each color channel so that the visual quality of the target image is consistent with the reference image input to the network. The image cropping is used to crop and scale the image to a uniform size to ensure that the size of the input image is consistent; Scale the image to a uniform size, specifically 256×256.

4. The underwater image enhancement method based on convolutional neural network according to claim 1, characterized in that: In step S2, the basic network architecture ResNet includes an input layer, a convolutional layer, a residual module, and a fully connected layer; The input layer is used to input an image, that is, the image of size 256×256 after the preprocessing; the convolution layer includes multiple convolution layers Conv and batch normalization BatchNorm layers, which are used to extract low-level features; the residual module includes multiple residual blocks, each residual block includes two convolution layers and a batch normalization layer; the fully connected layer is used to connect the above residual modules and output the final feature map.

5. The underwater image enhancement method based on convolutional neural network according to claim 1, characterized in that: In step S2, the self-attention mechanism is inserted into the middle layer of ResNet. The specific operation of the self-attention module is: S5-1, input feature map from the middle layer of ResNet, size is H×W×C; S5-2, perform transformation operations, generate a Query feature map through 1×1 convolution, generate a Key feature map through 1×1 convolution, and generate a Value feature map through 1×1 convolution; S5-3, similarity calculation, calculate the dot product similarity between Query and Key to obtain the attention weight matrix; S5-4, weighted summation, multiply the attention weight matrix by the Value feature map to obtain a weighted sum feature map; S5-5, output feature map, add the weighted sum feature map to the original input feature map to obtain an enhanced feature map.

6. The underwater image enhancement method based on convolutional neural network according to claim 1, characterized in that: In step S2, the feature pyramid network FPN is used for multi-scale feature fusion, and feature fusion is performed in two ways: bottom-up and top-down. Specifically, The bottom-up path of the FPN structure outputs feature maps of different scales from each stage of ResNet; The top-down path of the FPN structure starts from the highest-level feature map, gradually upsamples downward, and merges with the bottom-up feature map. The number of high-level feature map channels is halved through 1×1 convolution, and then added to the corresponding low-level feature map through 2×2 upsampling. The feature fusion is to further process the fused feature map through 3×3 convolution to obtain the final multi-scale feature map.

7. The underwater image enhancement method based on convolutional neural network according to claim 1, characterized in that: In step S3, the comprehensive loss function includes mean square error loss, structural similarity loss, perceptual loss and adversarial loss; The mean square error loss: Among them, MSE Loss represents the mean square error loss; y i Represents the i-th pixel value of the real image; Represents the i-th pixel value of the predicted image, 1≤i≤n; n represents the total number of pixels in the image; The structural similarity loss: in, represents the structural similarity loss; μ y and Represent the pixel averages of the real image and the predicted image respectively; σ y and Represent the variance of the pixel values ​​of the real image and the predicted image respectively; represents the covariance of the real image and the predicted image; C1 and C2 are structural constants; The perceived loss: Among them, Perceptual Loss Indicates perceived loss; express The L2 norm of , which represents the Euclidean distance; represents the j-th layer mapping of the training network; y and Represent the real image and the predicted image respectively; The adversarial loss: Adversarial Loss =E x～Pdata [logD(x)]+E z～Pz [log(1-D(G(z)))] Among them, Adversarial Loss represents the image adversarial loss; E x～Pdata represents the expectation of the real data x sampled from the real data distribution Pdata; E z～Pz represents the expectation of random noise z sampled from the noise distribution Pz; D and G represent the discriminator and generator in the training network respectively; logD(x) is the value of the discriminator's output on the real data x after being processed by the logarithmic function; log(1-D(G(z))) is the value of the discriminator's output on the generated data G(z) after being processed by the logarithmic function; The comprehensive loss function: Among them, Total Loss represents the comprehensive loss function; α, β, γ and δ are the weight parameters of mean square error loss, structural similarity loss, perceptual loss and adversarial loss, respectively.

8. The underwater image enhancement method based on convolutional neural network according to claim 1, characterized in that: In step S5, The PSNR objective evaluation index is calculated as follows: PSNR is based on mean square error (MSE) Loss The specific formula for calculation is: Among them, MSE Loss Represents mean square error; PSNR is an objective evaluation indicator; MAX I is the maximum pixel value of the image; The verification test is to test the trained model using the verification data set, generate enhanced images, and then calculate the PSNR objective evaluation index for each pair of original images and enhanced images; calculate the average PSNR value on the verification set; if it is lower than the set threshold, perform secondary training; Secondary training is to add the validation data set to the training data and retrain with different hyperparameters; retrain the model until the PSNR reaches the preset indicator; The threshold is determined by calculating the mean value, median and standard deviation of the error between the reference target image and the real image.

9. An underwater image enhancement device based on convolutional neural network, characterized in that: The device includes a data preprocessing module, a network architecture design module, a loss function design module, a model training module, a verification test module and a model optimization and deployment module; The data preprocessing module is used for image data collection, image normalization, image color correction and image cropping; it collects a variety of underwater image data, including actual captured images, AI synthesized images and public underwater image data sets; Normalize the pixel values ​​of the image to the range [-1,1] and adjust the gain of each color channel; crop and scale the image to a uniform size to ensure that the input image has the same size; The network architecture design module selects the ResNet architecture as the basic network architecture for initial feature extraction and image reconstruction, introduces a self-attention mechanism module in the middle layer of ResNet, and uses the feature pyramid network FPN for multi-scale feature fusion; The loss function design module is used to design a comprehensive loss function, which combines mean square error loss, structural similarity loss, perceptual loss and adversarial loss; The model training module is used to divide the image data into a training set, a validation set and a test set, and train the convolutional neural network model using a designed comprehensive loss function; The verification test module is used to verify the model performance, use the verification data set to test the trained model, calculate the average PSNR value on the verification set, and determine whether the model meets the preset performance standard. If the PSNR does not reach the set threshold, use the verification data set for secondary training; The model optimization and deployment module is used to optimize and deploy the trained model, and the model optimization includes model compression, model pruning and model quantization; Model deployment is to deploy the model to actual applications after model optimization is completed; The output end of the data preprocessing module is connected to the input end of the network architecture design module; The output end of the network architecture design module is connected to the input end of the loss function design module; the output end of the loss function design module is connected to the input end of the model training module; the output end of the model training module is connected to the input end of the verification and testing module; the output end of the verification and testing module is connected to the input end of the model optimization and deployment module.

Images