Underwater Image Enhancement Method Based on Knowledge Transfer

Abstract

This invention provides an underwater image enhancement method based on knowledge transfer. The method includes: constructing an underwater image enhancement network model; building a first loss function; constraining the first network model using the first loss function, and training the first network model with N underwater images to obtain a teacher network model; constructing a second loss function based on the teacher network model; constraining the second network model using the second loss function, and training the second network model with N underwater images and M transformed images to obtain a student network model; inputting the original image into the student network model, and the student network model performing enhancement processing on the original image. By constraining the training of the second network model with the second loss function, the generated image is made closer to the reference image, with more vivid colors, higher contrast, and richer texture and detail.

Description

Technical Field

[0001] This invention relates to the field of underwater image enhancement technology, and more particularly to an underwater image enhancement method based on knowledge transformation. Background Technology

[0002] The unknown and mysterious underwater environment and its precious underwater resources have become a hot research topic. However, light suffers severe attenuation and scattering when propagating in water, leading to blue-green tint and low contrast in underwater images. Degraded underwater images significantly affect the efficiency of autonomous underwater vehicles (AUVs and ROVs), further limiting ocean exploration. Therefore, developing effective underwater image enhancement methods to restore the color and contrast of underwater images is of great significance to various fields such as marine biology, oceanography, and underwater archaeology.

[0003] Currently, various underwater image enhancement methods have been proposed and have achieved certain results, but there are also many problems and significant room for improvement. Existing underwater image enhancement methods can be divided into traditional underwater image enhancement methods and deep learning-based underwater image enhancement methods. Traditional underwater image enhancement methods can be further divided into non-physical model methods and physical model-based methods.

[0004] Non-physical model underwater image enhancement methods typically process the pixel distribution of the acquired underwater image directly, adjusting the color distribution through operations such as stretching and blending to improve color cast and enhance contrast. While traditional non-physical model underwater image enhancement methods can quickly improve underwater image quality, most lack specific reference standards, leading to inconsistencies in subjective and objective evaluations of the restored underwater images. Furthermore, the restored underwater images often exhibit under-enhancement or over-exposure.

[0005] On the other hand, physical model-based methods generally rely on specific physical imaging models. These methods treat underwater image attenuation as a reversible process, thus simulating the attenuation process by establishing a physical model. Then, based on assumptions or prior estimates, the parameters in the constructed imaging model are estimated, and finally, the image before attenuation is solved in reverse. While traditional physical model-based underwater image enhancement methods can recover underwater images before degradation to some extent, their low generalization ability and the resulting variable underwater images are due to the complexity of the constructed imaging models, the large number of parameters, and the difficulty in accurately estimating them.

[0006] Finally, with the improvement of computing hardware performance and the abundance of data samples, deep learning technology has been widely applied in the field of image processing. In the underwater environment, deep learning-based underwater image enhancement methods have also shown promising performance. For distorted underwater images, deep learning-based methods first extract feature information from the underwater image by constructing a network model, and then establish a mapping between low-quality and high-quality underwater images under the supervision of a large number of data samples and various loss functions, thereby completing the underwater image enhancement task. Similarly, underwater images generated by deep learning-based underwater image enhancement methods are prone to a reddish tint, because most of these methods do not specifically address the color degradation of underwater images.

[0007] Therefore, there is a need to provide an underwater image enhancement method based on knowledge conversion. Summary of the Invention

[0008] To address the issue of reddish tint in existing underwater image enhancement techniques, this invention provides a knowledge-based underwater image enhancement method. This invention comprehensively considers the relationship between the red, green, and blue channels to correct underwater image colors, alleviating the reddish tint problem. Simultaneously, it uses knowledge conversion technology to enrich the dataset required for model training, supervising the model's training and thereby improving the generalization ability of the knowledge-based underwater image enhancement method provided by this invention. This allows the method to better adapt to diverse underwater scenarios and enhance the overall underwater image enhancement effect.

[0009] The technical means employed in this invention are as follows:

[0010] This invention provides an underwater image enhancement method based on knowledge conversion, comprising:

[0011] Construct a dataset, which includes a first training set and a second training set. The first training set includes N underwater images and N corresponding reference images. The second training set includes M synthetic underwater images, M transformed images corresponding to the M synthetic underwater images, and M reference images corresponding to the M synthetic underwater images.

[0012] An underwater image enhancement network model is constructed, which includes a first network model and a second network model with identical structures.

[0013] Construct the first loss function;

[0014] The first network model is constrained by the first loss function, and the first network model is trained using N underwater images to obtain a teacher network model;

[0015] The second loss function is constructed based on the teacher network model, and the second loss function is calculated according to the following formula:

[0016]

[0017] Among them, λ1=0.5, λ2=0.2, λ3=0.05, λ4=0.1, Let the mean absolute error loss function be . For the perceptual loss function, Let be the gradient loss function. The knowledge transfer loss function;

[0018] The second network model is constrained by the second loss function, and the student network model is obtained by training the second network model with N underwater images and M transformed images;

[0019] The original image is input into the student network model, and the student network model performs enhancement processing on the original image.

[0020] Furthermore, the acquisition of the transformed image includes:

[0021] Select P real underwater images and P composite underwater images;

[0022] The CycleGAN network is trained using P real underwater images and P synthesized underwater images.

[0023] The M synthetic underwater images are input into the trained CycleGAN network to obtain M transformed images.

[0024] Furthermore, the underwater image enhancement network model includes a weight parameter estimation subnetwork, a first color correction subnetwork, a second color correction subnetwork, and a third color correction subnetwork;

[0025] The first color correction sub-network performs color compensation for the red channel according to the following formula:

[0026]

[0027] The second color correction sub-network performs color compensation for the green channel according to the following formula:

[0028]

[0029] The third color correction sub-network performs color compensation for the blue channel according to the following formula:

[0030]

[0031]

[0032] Where x is a pixel, I λC (x) represents the result after color compensation, λ∈{R,G,B},I λ (x) represents a single color channel of the input image. This is a graph of weight parameters obtained from the weight parameter estimation subnetwork.

[0033] Furthermore, the weight parameter estimation subnetwork includes a first convolutional block to a sixth convolutional block, the first to fifth convolutional blocks including 3×3 convolutional layers and LeakyReLU activation functions, and the sixth convolutional block including 3×3 convolutional layers and Tanh activation functions.

[0034] Furthermore, the first loss function is calculated in the following manner:

[0035]

[0036] Among them, E X，Y To calculate the expected value of the absolute difference, (i, j) represents the pixel coordinates, G(X(i, j)) represents the generated image, Y(i, j) represents the reference image, and Φ... j The feature map is obtained by the j-th convolution of a VGG-19 neural network model pre-trained on the ImageNet dataset. The horizontal gradient of the reference image. The vertical gradient of the reference image.

[0037] Furthermore, the knowledge transfer loss function is calculated according to the following formula:

[0038]

[0039] Among them, G t G(X) is the image generated by the teacher model, and G(X) is the image generated by the student model.

[0040] Furthermore, the construction of the underwater image enhancement network model includes:

[0041] The underwater image enhancement network model is configured to use an image processing unit to process the input image in order to train the underwater image enhancement network model;

[0042] The underwater image enhancement network model is optimized using the Adam algorithm.

[0043] The learning efficiency of the underwater image enhancement network model is set to 0.0001;

[0044] The input image size is set to 256×256×3, and the pixel values ​​are normalized to [-1,1].

[0045] The batch size is set to 8.

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

[0047] 1. The underwater image enhancement network model provided by this invention considers the relationship between the red, green and blue color channels of underwater images, and makes full use of the residual information of each color channel to perform color correction on the other two color channels, effectively alleviating the phenomenon of excessive red in the generated image and correcting the color of the underwater image better.

[0048] 2. The underwater image enhancement method based on knowledge conversion provided by this invention guides the training of the student model through the teacher model and constrains the training of the second network model through the second loss function, which effectively improves the generalization ability of the student model and promotes the generated image to be closer to the reference image, with more vivid colors, higher contrast, and richer texture and details.

[0049] 3. The underwater image enhancement method based on knowledge conversion provided by this invention is a lightweight network with a simple model, few parameters, and fast computation speed. Attached Figure Description

[0050] To more clearly illustrate the technical solutions in the embodiments of the present invention or the prior art, the drawings used in the description of the embodiments or the prior art will be briefly introduced below. Obviously, the drawings described below are some embodiments of the present invention. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.

[0051] Figure 1 A flowchart of the underwater image enhancement method based on knowledge conversion provided by the present invention.

[0052] Figure 2 This is a diagram of the network model structure.

[0053] Figure 3 This is a structural diagram of the residual network layers used in the network model.

[0054] Figure 4 This is a subjective evaluation diagram of the method of the present invention and the comparative method.

[0055] Figure 5 This is another subjective evaluation diagram of the method of the present invention and the comparative method.

[0056] Figure 6 This is another subjective evaluation diagram of the method of the present invention and the comparison method. Detailed Implementation

[0057] To enable those skilled in the art to better understand the present invention, the technical solutions of the present invention will be clearly and completely described below with reference to the accompanying drawings of the embodiments of the present invention. Obviously, the described embodiments are only some embodiments of the present invention, and not all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort should fall within the scope of protection of the present invention.

[0058] It should be noted that the terms "first," "second," etc., in the specification, claims, and accompanying drawings of this invention are used to distinguish similar objects and are not necessarily used to describe a specific order or sequence. It should be understood that such data can be interchanged where appropriate so that the embodiments of the invention described herein can be implemented in orders other than those illustrated or described herein. Furthermore, the terms "comprising" and "having," and any variations thereof, are intended to cover a non-exclusive inclusion; for example, a process, method, system, product, or apparatus that comprises a series of steps or units is not necessarily limited to those steps or units explicitly listed, but may include other steps or units not explicitly listed or inherent to such processes, methods, products, or apparatus.

[0059] Reference Figure 1 , Figure 2 and Figure 3 , Figure 1 The flowchart below shows the underwater image enhancement method based on knowledge conversion provided by this invention. Figure 2 This is a diagram of the network model structure. Figure 3 The diagram shows the structure of the residual network layers used in the network model, illustrating a specific embodiment of the underwater image enhancement method based on knowledge conversion provided by this invention, including:

[0060] Construct a dataset, which includes a first training set and a second training set. The first training set includes N underwater images and N corresponding reference images. The second training set includes M synthetic underwater images, M transformed images corresponding to the M synthetic underwater images, and M reference images corresponding to the M synthetic underwater images.

[0061] Among them, reference Figure 1 The acquisition of the transformed images includes: selecting P real underwater images and P synthesized underwater images; training the CycleGAN network using the P real underwater images and P synthesized underwater images; and inputting M synthesized underwater images into the trained CycleGAN network to obtain M transformed images.

[0062] Understandably, the P real underwater images can be 1000 severely degraded real underwater images selected from the real dataset ChinaMM. The P synthetic underwater images can be 1000 synthetic underwater images containing various styles selected from the synthetic dataset UWCNN. The M synthetic underwater images can be 1000 synthetic underwater images containing various styles selected from the synthetic dataset UWCNN. Guided by real underwater images, the synthetic underwater images from the synthetic dataset UWCNN are transformed into transformed images that are more severely degraded and closer to the real underwater environment. The first training set includes 800 pairs of underwater images selected from the dataset UIEB and 1000 pairs of underwater images selected from the dataset EUVP. There is no overlap between the datasets, but it is not limited to these.

[0063] An underwater image enhancement network model is constructed, which includes a first network model and a second network model with the same structure.

[0064] Among them, reference Figure 2 and Figure 3 The underwater image enhancement network model includes a weight parameter estimation subnetwork, a first color correction subnetwork, a second color correction subnetwork, and a third color correction subnetwork.

[0065] The first color correction sub-network performs color compensation for the red channel according to the following formula:

[0066]

[0067] The second color correction sub-network performs color compensation for the green channel according to the following formula:

[0068]

[0069] The third color correction sub-network performs color compensation for the blue channel according to the following formula:

[0070]

[0071] Where x is a pixel, I λC (x) represents the result after color compensation, λ∈{R,G,B},I λ (x) represents a single color channel of the input image. This is a graph of the weight parameters obtained from the weight parameter estimation subnetwork.

[0072] The weight parameter estimation subnetwork includes the first to the sixth convolutional blocks. The first to the fifth convolutional blocks each contain a 3×3 convolutional layer and a LeakyReLU activation function, while the sixth convolutional block contains a 3×3 convolutional layer and a Tanh activation function.

[0073] Understandably, the weight parameter estimation subnetwork is used to obtain multiple weight parameter maps to control the degree of underwater image enhancement and mitigate under-enhancement and over-enhancement phenomena. Specifically, the first to fifth convolutional blocks of the weight parameter estimation subnetwork are used to extract underwater image features, and the sixth convolutional block is used to generate six weight parameter maps. (Refer to...) Figure 2 The weight parameter estimation subnetwork takes a 256×256×3 degraded underwater image as input. The first to fifth convolutional blocks output feature maps of sizes 256×256×8, 256×256×16, 256×256×32, and 256×256×32, respectively. The sixth convolutional block generates a weight parameter map of size 256×256×6. The first, second, and third color correction subnetworks compensate for the red, green, and blue channels of the underwater image, respectively. These subnetworks use the same design structure, restoring the color of their respective channels based on the residual information from the other two color channels. To reduce the number of parameters in the underwater image enhancement network model and improve algorithm efficiency, this embodiment uses only 6 residual blocks and a 1×1 convolution operation to implement the color correction process. Figure 3 As shown, each residual block consists of two 3×3 convolution operations, a block normalization operation, and a ReLU activation function. The concise underwater image enhancement network model design enables the underwater image enhancement method based on knowledge conversion proposed in this invention to have high image processing speed and real-time underwater image enhancement capabilities.

[0074] Construct the first loss function;

[0075] The first loss function is calculated as follows:

[0076]

[0077] Among them, E X，Y To calculate the expected value of the absolute difference, (i, j) represents the pixel coordinates, G(X(i, j)) represents the generated image, Y(i, j) represents the reference image, and Φ... j To extract the feature map obtained through the j-th convolution of the neural network model, this implementation uses a VGG-19 network pre-trained on the ImageNet dataset as the feature extraction network, and calculates the perceptual loss using the output features of the ReLU5-2 layer of this network. To reference the horizontal gradient of the image, This refers to the vertical gradient of the reference image. Let the mean absolute error loss function be . For the perceptual loss function, The gradient loss function and the mean absolute error loss function can constrain the generated image at the pixel level, making the generated image closer to the reference image in terms of color, etc.; the perceptual loss function constrains the generated image at the image feature level, making the generated image closer to the reference image in terms of texture, detail, etc.; the gradient loss function can improve the similarity of the edge features between the generated image and the reference image, thereby improving the contrast of the image and reducing the appearance of halos.

[0078] The first network model is constrained by the first loss function, and the teacher network model is obtained by training the first network model with N underwater images.

[0079] A second loss function is constructed based on the teacher network model, and the second loss function is calculated according to the following formula:

[0080]

[0081] Where λ1 = 0.5, λ2 = 0.2, λ3 = 0.05, and λ4 = 0.1 are all empirical weighting numerators for the equilibrium loss components. Let the mean absolute error loss function be . For the perceptual loss function, Let be the gradient loss function. The knowledge transfer loss function;

[0082] The knowledge transfer loss function is calculated using the following formula:

[0083]

[0084] Among them, G t G(X) is the image generated by the teacher model, and G(X) is the image generated by the student model.

[0085] Understandable, The loss function is a knowledge transfer loss function built on the mean absolute error loss, which promotes the transfer of knowledge from the teacher model to the student model and improves the generalization ability of this method.

[0086] The second network model is constrained by a second loss function, and the student network model is obtained by training the second network model with N underwater images and M transformed images.

[0087] The original image is input into the student network model, which then performs enhancement processing on the original image.

[0088] In some optional embodiments, an underwater image enhancement network model is constructed, including:

[0089] The underwater image enhancement network model is trained by processing the input image using an image processing unit; the Adam algorithm is used for optimization; the learning efficiency of the underwater image enhancement network model is set to 0.0001; the input image size is set to 256×256×3, and the pixel values ​​are normalized to [-1,1]; the BatchSize is set to 8.

[0090] In some alternative embodiments, refer to Figure 4 , Figure 5 and Figure 6 , Figure 4 This is a subjective evaluation diagram of the method of the present invention and the comparison method. Figure 5 This is another subjective evaluation diagram comparing the method of the present invention with the comparative method. Figure 6 This is another subjective evaluation graph for the method and comparison method of the present invention. The dataset also includes a test set, which includes 500 real underwater images selected from the real dataset ChinaMM, 90 pairs of underwater images selected from the dataset UIEB, and 1000 pairs of underwater images selected from the dataset EUVP. The test set has no overlap with the first training set and the second training set.

[0091] After obtaining the student model, the image enhancement effect of the student model must be evaluated subjectively and objectively based on the test set, including:

[0092] A comparative experiment was conducted using 10 of the most representative underwater image enhancement methods in the existing technology and the knowledge-transformation-based underwater image enhancement method provided in this embodiment. The selected comparison methods are: UDCP, UIBLA, ColorBalance, GDCP, Water-net, FUnIE, UWCNN, and CWR. Figures 4 to 6 Using UCCNet as the teacher model and UCCNet-KT as the student model, the knowledge-transformation-based underwater image enhancement method provided in this embodiment exhibits excellent performance for any given underwater image. Other methods show significant differences in enhancement effects for different underwater images, indicating that the knowledge-transformation-based underwater image enhancement method provided in this embodiment has higher generalization ability and robustness. Furthermore, the knowledge-transformation-based underwater image enhancement method provided in this embodiment can effectively improve the contrast of underwater images, reduce color cast, and enhance texture details. Especially for underwater images from the EUVP and UIEB datasets, the underwater images enhanced by the knowledge-transformation-based underwater image enhancement method provided in this embodiment are closer to the reference images.

[0093] The above is a subjective evaluation, and it also includes objective evaluations:

[0094] Eight objective evaluation metrics were selected to quantitatively analyze the results of the underwater image enhancement method based on knowledge conversion provided in this embodiment and existing underwater image enhancement methods. For underwater images with reference images, this embodiment used four reference-based evaluation metrics—mean squared error (MSE), structural similarity (SSIM), peak signal-to-noise ratio (PSNR), and PCQI—to evaluate the results of each method. A smaller MSE value indicates that the enhanced underwater image is closer to the reference image, while larger SSIM, PSNR, and PCQI values ​​indicate that the enhanced underwater image is more similar to the reference image in terms of structure, detail, and contrast. For underwater images without reference images, four non-reference-based evaluation metrics—UIQM, UCIQE, Natural Image Quality (NIQE), and Entropy—were used to evaluate the results of each method. A smaller NIQE value indicates that the enhanced underwater image is closer to a natural image taken in the air, while larger UIQM, UCIQE, and Entropy values ​​indicate that the generated image has richer colors, higher contrast, and is more aesthetically pleasing. Referring to Tables 1 to 3, the quantitative analysis of both the teacher model UCCNet and the student model UCCNet-KT shows significant advantages. This indicates that the underwater image enhancement method based on knowledge transfer provided in this embodiment can effectively enhance degraded underwater images, generating images with rich colors, high contrast, and aesthetic appeal, which is consistent with the conclusions of subjective evaluation. Furthermore, Tables 1-3 also show that the student model UCCNet-KT possesses stronger generalization ability compared to the teacher model UCCNet, further demonstrating the effectiveness of knowledge transfer.

[0095] Table 1. Quantitative comparisons on the EUVP test set

[0096]

[0097] Table 2. Quantitative Comparison Table on the UIEB Test Set

[0098]

[0099] Table 3. Quantitative Comparison Table on the ChinaMM Test Set

[0100]

[0101] In summary, the underwater image enhancement method based on knowledge conversion provided in this embodiment can effectively enhance degraded underwater images, alleviate under-enhancement or over-enhancement phenomena during the enhancement process, and generate images with vivid colors, clear textures, and high contrast. Furthermore, the underwater image enhancement method based on knowledge conversion provided in this embodiment also possesses high generalization ability and robustness, and can adapt well to diverse and complex underwater environments.

[0102] The sequence numbers of the above embodiments of the present invention are for descriptive purposes only and do not represent the superiority or inferiority of the embodiments.

[0103] In the above embodiments of the present invention, the descriptions of each embodiment have different focuses. For parts not described in detail in a certain embodiment, please refer to the relevant descriptions of other embodiments.

[0104] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention, and not to limit them; although the present invention has been described in detail with reference to the foregoing embodiments, those skilled in the art should understand that modifications can still be made to the technical solutions described in the foregoing embodiments, or equivalent substitutions can be made to some or all of the technical features; and these modifications or substitutions do not cause the essence of the corresponding technical solutions to deviate from the scope of the technical solutions of the embodiments of the present invention.

Claims

1. An underwater image enhancement method based on knowledge transfer, characterized in that, The method comprises the following steps: constructing a data set comprising a first training set and a second training set, the first training set comprising N underwater images and corresponding N reference images, and the second training set comprising M synthetic underwater images, M transformed images corresponding to the M synthetic underwater images, and M reference images corresponding to the M synthetic underwater images; building an underwater image enhancement network model, the underwater image enhancement network model comprising a first network model and a second network model with the same structure; constructing a first loss function; using the first loss function to constrain the first network model, and using the N underwater images to train the first network model to obtain a teacher network model; constructing a second loss function according to the teacher network model, the second loss function being calculated according to the following formula: ； wherein, = 0.5, = 0.2, = 0.05, = 0.1, is the mean absolute error loss function, is the perceptual loss function, is the gradient loss function, is the knowledge transfer loss function; the knowledge conversion loss function being calculated according to the following formula: ； wherein, images generated for the teacher model, images generated for the student model; using the second loss function to constrain the second network model, and using the N underwater images and the M transformed images to train the second network model to obtain a student network model; inputting an original image into the student network model, and the student network model performing enhancement processing on the original image.

2. The knowledge conversion based underwater image enhancement method according to claim 1, wherein, The acquisition of the transformed image comprises the following steps: selecting P real underwater images and P synthetic underwater images; training a CycleGAN network using the P real underwater images and the P synthetic underwater images; inputting the M synthetic underwater images into the trained CycleGAN network respectively to obtain M transformed images.

3. The knowledge conversion based underwater image enhancement method according to claim 1, wherein, The underwater image enhancement network model comprises a weight parameter estimation sub-network, a first color correction sub-network, a second color correction sub-network, and a third color correction sub-network. The first color correction sub-network performs red channel color compensation according to the following formula: ； The second color correction sub-network performs green channel color compensation according to the following formula: ； The third color correction sub-network performs blue channel color compensation according to the following formula: ； wherein, is a pixel point, is a result of color compensation, , is a single color channel of the input image, is a weight parameter map obtained by the weight parameter estimation subnetwork.

4. The knowledge conversion based underwater image enhancement method according to claim 3, wherein, The weight parameter estimation sub-network comprises a first convolutional block to a sixth convolutional block, the first convolutional block to the fifth convolutional block comprising a 3×3 convolutional layer and a LeakyReLU activation function, and the sixth convolutional block comprising a 3×3 convolutional layer and a Tanh activation function.

5. The knowledge conversion based underwater image enhancement method according to claim 1, wherein, The first loss function is calculated according to the following formula: ； wherein, is the expectation value of the absolute difference, is the pixel coordinate, is the generated image, is the reference image, is the feature map obtained by the 1st convolution of the VGG-19 neural network model pre-trained on the ImageNet dataset, is the feature map obtained by the 2nd convolution of the VGG-19 neural network model pre-trained on the ImageNet dataset, , is the horizontal gradient of the reference image, is the vertical gradient of the reference image.

6. The knowledge conversion based underwater image enhancement method according to claim 1, wherein, The building of the underwater image enhancement network model comprises the following steps: setting the underwater image enhancement network model to use an image processing unit to process an input image to train the underwater image enhancement network model; setting the underwater image enhancement network model to use an Adam algorithm for optimization; setting the learning efficiency of the underwater image enhancement network model to 0.0001; setting the size of the input image to 256×256×3, and normalizing the pixel value to [-1, 1]; BatchSize is set to 8.

Images