An underwater image restoration pre-training-fine-tuning method and system based on semantic and diffusion double priors and a medium

Abstract

The application discloses a kind of underwater image restoration pre-training-fine-tuning method, system and medium based on semantic and diffusion double prior, specific method: the basic pre-training network of shallow / deep feature extraction and double task reconstruction module is constructed;Adopting multiple degradation mode blind degradation strategy constructs pre-training dataset, trains model to extract generalization visual feature;Through embedding adapter module, fine-tune model on UIEB dataset with gradual thaw strategy, learn underwater specific illumination attenuation, color deviation distribution and scattering effect, realize cross-domain knowledge transfer;For denoising and super-resolution task embedding semantic prior module, fuse the semantic mask generated by SAM to strengthen image structure understanding;For super-resolution task, introduce diffusion prior module, play the advantage of diffusion model in complex image generation task, provide fine noise modeling and detail reconstruction capability for underwater image restoration model.The application has the advantages of maintaining color balance, restoring image details real delicate and the like.

Description

TECHNICAL FIELD

[0001] The application belongs to the field of image processing, and particularly relates to an underwater image restoration pre-training-fine-tuning method and system based on semantic and diffusion double priors and a medium. BACKGROUND

[0002] In marine environment detection and research, underwater images are key data carriers for obtaining the dynamic of ecological system, the characteristics of geological structure and the state of underwater engineering, and the quality of the underwater images directly affects the reliability of scientific research conclusions and the accuracy of engineering decisions. However, in the actual imaging process, the interaction of photons with water molecules and suspended particles leads to scattering phenomena (including Rayleigh scattering and Mie scattering), dissolved substances and suspended particles selectively absorb light of different wavelengths, and refraction at the medium interface will cause light deflection. These factors together cause the image to present serious quality degradation. Specifically, the phenomena include blurred outlines of distant targets, lost details of close-up scenes, color shift and contrast attenuation, which restricts the actual application effect of underwater visual analysis technology.

[0003] For such underwater image restoration tasks, researchers have proposed two technical routes: the first type of method is based on optical physics theory, and improves the image quality by constructing a light propagation attenuation model, combining scattering coefficient estimation and energy compensation. The second type of method uses the nonlinear feature mapping ability of deep neural networks to train an end-to-end image restoration model with a paired data set. The former has a theoretical basis, but it is difficult to describe complex environmental factors such as water turbulence and uneven distribution of suspended matter, and real-time estimation of key parameters such as scattering coefficients is also a great challenge. These inherent limitations have gradually shifted the focus of research to deep learning methods, hoping to use the powerful feature extraction and nonlinear modeling capabilities of neural networks to obtain better restoration results while ensuring computational efficiency.

[0004] Underwater image restoration methods based on deep learning can be mainly divided into two categories: CNN-based methods and Transformer-based methods, according to the different structures of backbone networks. The design of convolutional neural networks is inspired by the neural connection patterns in the visual cortex of animals, where a single neuron responds to local area stimulation, and the receptive fields of different neurons partially overlap to effectively cover the entire visual field. Convolutional neural networks are composed of multiple convolutional layers, pooling layers, and fully connected layers at the top, which process two-dimensional image data by reducing the data dimension and enhancing the robustness of the network to efficiently extract spatial features and local information. In recent years, CNN-based underwater image restoration methods have made significant progress. For example, in 2019, Wang N, Zhou Y, Han F, et al. UWGAN: Underwater GAN for real-world underwater color restoration and dehazing [J]. arXiv preprint arXiv:1912.10269, 2019. used the UNet architecture for underwater image restoration, which captures low-level features while restoring high-level details through an encoder-decoder design, making it suitable for handling complex optical degradation phenomena in underwater images. However, CNN-based methods generally have limited local perception and insufficient modeling of long-range dependencies. In contrast, Transformer models exhibit stronger advantages in capturing global information and long-range dependencies. Through self-attention mechanisms, Transformers can more comprehensively understand the relationships between regions in images, especially excelling at handling global structural information in images. Transformer-based underwater image restoration methods can effectively address issues such as long-distance object blurring and color distortion in underwater imaging. In 2023, Khan M D R, Kulkarni A, Phutke S S, et al. Underwater image enhancement with phase transfer and attention [C] / / 2023 International Joint Conference on Neural Networks (IJCNN). IEEE, 2023:1-8. proposed a lightweight phase-based Transformer network (Phase-Transformer) for underwater image restoration. This method focuses on using phase-based self-attention mechanisms to extract uncontaminated image features and introduces an optimized phase attention module to restore structural information in the image. The phase attention module learns the attention features of the input image to accurately restore image details, demonstrating its potential applications in the field of underwater image restoration.However, the Transformer-based method generally has high computational complexity. In addition, some researchers attempt to combine the underwater image restoration model with the diffusion model. For example, Nathan O B, Levy D, Treibitz T, et al. Osmosis: Rgbd diffusion prior for underwater image restoration [C] / / European Conference on Computer Vision. Cham: Springer Nature Switzerland, 2024: 302-319. restores underwater images by using the RGBD diffusion prior. However, the introduction of the diffusion model generally causes the inference time to be too long.

[0005] In summary, the underwater image restoration method based on deep learning can automatically learn and model the complex degradation process of underwater images. However, the CNN-based method generally has the problems of limited local perception and insufficient modeling of long-distance dependency relationship. The Transformer-based method generally has the problem of high computational complexity. In addition, there is an important dilemma in the underwater image restoration task of the deep learning method: the lack of real high-quality underwater images. The current solutions mainly rely on two pairing methods: one is to pair real low-quality underwater images with artificially simulated high-quality images; and the other is to pair simulated underwater degraded images with real non-underwater high-quality images. Both of the two schemes have obvious limitations. SUMMARY

[0006] The underwater image restoration pre-training-fine-tuning method based on the semantic and diffusion double prior, system and medium provided by the present application can better improve the quality of the restored underwater images.

[0007] The object of the present application is achieved by the following technical solutions:

[0008] The underwater image restoration pre-training-fine-tuning method based on the semantic and diffusion double prior, system and medium provided by the present application can better improve the quality of the restored underwater images.

[0009] Step 1: a blind degradation strategy is used to construct a pre-training data set for pre-training; a basic pre-training network including a shallow feature extraction module, a deep feature extraction module and an image reconstruction module is constructed; a universal high-quality original image set is obtained and is subjected to a blind degradation operation to generate a training data set, and the basic pre-training network is trained by using the training data set and the universal high-quality original image set; the image reconstruction module includes parallel denoising tasks and super-resolution tasks;

[0010] Step 2: Optimize the deep feature extraction module on the underwater image dataset using the adapter-fine-tuning method; the adapter-fine-tuning method embeds an adapter in the deep feature extraction module and uses the collected low-quality underwater image dataset I LQ_water The parameters of the adapter are fine-tuned through the multi-stage fine-tuning strategy of progressive unfreezing, which optimizes the deep feature extraction module in the basic pre-trained network trained in step 1;

[0011] Step 3: Introduce a semantic prior module for the denoising task and the super-resolution task; embed a semantic prior module after the adapter, use the SAM model to generate a corresponding semantic segmentation mask M for each image in the dataset I LQ_water , and fuse the deep features extracted by the deep feature extraction module optimized in step 2, the semantic segmentation mask M, and the SAM prior representation through the SPT Unit to obtain the feature map F i n , and use F i n to obtain the final high-quality denoised underwater image result I HQ_wdn and the super-resolution high-quality underwater image I HQ1_wsr2 , I HQ1_wsr4 ;

[0012] Step 4: Further introduce a diffusion prior module for the super-resolution task; use a pre-trained diffusion model G θ composed of a trainable VAE encoder E θ , a fine-tuned diffusion network ε θ , and a fixed VAE decoder D θ ; introduce a trainable LoRA layer in E θ and ε θ , and fine-tune the parameters of the model G θ using the collected underwater image dataset; input I HQ1_wsr2 , I HQ1_wsr4 to the fine-tuned model G θ , generate image embeddings through the text prompt extractor DAPE, drive single-step latent space denoising, and generate the final super-resolution high-quality image I HQ2_wsr2 , I HQ2_wsr4 .

[0013] Further, the universal high-quality original image set in step 1 is subjected to blind degradation operation, including three core parts: blur degradation, scale transformation, and noise simulation, and a second-order degradation mechanism is introduced to double superimpose the "blur-noise" degradation process.

[0014] Further, the training dataset is first subjected to a shallow feature extraction module to obtain shallow features F shallow, and then a deep feature extraction module composed of stacked residual blocks is used to obtain a deep feature F deep , and then a denoising task of an image reconstruction module is performed

[0015] I HQ_dn LQ +Conv(F deep )

[0016] wherein I HQ_dn represents a reconstructed high-quality denoised image, Conv represents a single-layer convolution operation with a 3*3 convolution kernel, represents a training data set;

[0017] The super-resolution task adopts a progressive nearest neighbor interpolation up-sampling, and performs a two-fold super-resolution as follows:

[0018] F1=LeakyReLU(Conv(F deep ))

[0019] F2=LeakyReLU(Conv(Nearest(F1)))

[0020] I HQ_2SR =Conv(LeakyReLU(Conv(F2)))

[0021] On the basis of the two-fold super-resolution image reconstruction, a four-fold super-resolution image reconstruction is performed as follows:

[0022] F3=LeakyReLU(Conv(Nearest(F2)))

[0023] I HQ_4SR =Conv(LeakyReLU(Conv(F3)))

[0024] wherein F1, F2 and F3 represent intermediate features; I HQ_2SR , I HQ_4SR are two-fold and four-fold super-resolution high-quality images, respectively; Nearest represents nearest neighbor interpolation; LeakyReLU represents a leaky rectified linear unit;

[0025] Each residual block is composed of a single-layer MambaVision.

[0026] ​Further, the step 2 inserts an adapter after each residual block; the multi-stage fine-tuning strategy of progressive un-freezing is: in the denoising task, it is divided into three stages: the first stage optimizes the adapter parameters, this stage un-freezes the image reconstruction module and the deep feature projection layer, and the last stage un-freezes all parameters; in the super-resolution task, it is divided into two stages: the first stage optimizes the adapter, the reconstruction module and the projection layer, and the last stage un-freezes all parameters; in all stages, a decreasing learning rate strategy is used to ensure stable convergence of the model.

[0027] Further, the SPTUnit fusion in step 3 is:

[0028] Generate semantic segmentation mask M = SAM(I LQ_water ); the SAM prior of the initial stage is represented as P0 = Conv(ReLU(Conv(Concat(I LQ_water , M))), where Concat represents a connection operation in the channel dimension, and ReLU represents an activation function; the output of the i-th SPTUnit is: i n = F i + αF i spt , where F i represents the feature map output by the i-th residual block in the deep feature extraction module, α represents a constant, and F i spt = Conv(ReLU(Conv(F i ')))*Conv(ReLU(Conv(P i )))+F i ', where P i represents the SAM prior F i ' that the i-th SPTUnit needs, and F T ' represents the intermediate representation generated by the i-th SPTUnit.

[0029] Further, the single-step latent space denoising in step 4 is:

[0030]

[0031] where α T and β y represent scaling factors at the T-th diffusion time step, c L is a text embedding, and z HQ1_wsr2 is a latent representation of the input image at the T-th diffusion time step, and the text prompt extractor DAPE is used to process different resolution input images I HQ1_wsr4 , I y1 to obtain text embeddings c HQ1_wsr2 , respectively.), c y2 = DAPE(I HQ1_wsr4 ), the final super-resolution high-quality underwater image I HQ2_wsr2 , I HQ2_wsr4 :

[0032]

[0033] Further, in steps 1 to 3, the loss function is an MSE loss, and in step 4, the loss function is a combination of MSE loss, LPIPS loss and Content loss.

[0034] A computer device / apparatus / system, comprising a memory, a processor and a computer program stored on the memory, the processor executing the computer program to implement the steps of the underwater image restoration pre-training-fine-tuning method based on semantic and diffusion dual priors.

[0035] A computer readable storage medium having stored thereon a computer program / instructions, which, when executed by a processor, implements the steps of the underwater image restoration pre-training-fine-tuning method based on semantic and diffusion dual priors.

[0036] The beneficial effects of the present application are:

[0037] (1) To solve the problem of training data scarcity in underwater image restoration tasks, a pre-training-fine-tuning model for underwater image restoration based on MambaVision is proposed. First, a blind degradation strategy is used to construct a pre-training dataset for pre-training. The blind degradation process simulates multiple degradation modes of different degrees, enabling the model to adapt to the feature distribution of complex environments and thus extracting more general and robust visual feature representations. On this basis, an adapter-fine-tuning method is used to optimize the model on the underwater image dataset UIEB, enabling the model to learn the unique degradation patterns in underwater environments, such as illumination attenuation, color bias distribution and scattering effects, and achieving efficient transfer of cross-domain knowledge to more accurately restore the color and details of underwater images. Experimental results on the underwater image dataset UIEB show that this method outperforms underwater image restoration models such as UWGAN and WaterNet in terms of PSNR, SSIM and other indicators;

[0038] (2) To solve the problems of semantic information blurring and detail loss in underwater image restoration, an underwater image restoration method based on semantic and diffusion double priors is proposed. In the pre-training-fine-tuning model of underwater image restoration based on MambaVision, the prior information of the semantic segmentation model (SegmentAnything Model, SAM) and the diffusion model (Stable Diffusion) is embedded. First, for the denoising task and super-resolution task, the semantic prior module is introduced to strengthen the understanding of the semantic boundary and target structure of the image. On this basis, for the super-resolution task, the diffusion prior module is further introduced to take advantage of the diffusion model in complex image generation tasks, providing fine noise modeling and detail reconstruction capabilities for the underwater image restoration model. By combining the semantic prior module and the diffusion prior module, the model can capture global semantic features and reconstruct fine-grained visual details. Experiments on the underwater image dataset UIEB show the effectiveness of the method. BRIEF DESCRIPTION OF DRAWINGS

[0039] Figure 1 is a flowchart of the method of the present invention;

[0040] Figure 2 is a schematic diagram of the overall structure of the method of the present invention;

[0041] Figure 3 is a comparison chart of the denoising effects of different models on the UIEB dataset;

[0042] Figure 4 is a comparison chart of the two-fold super-resolution processing effects of different models on the UIEB dataset;

[0043] Figure 5 is a comparison chart of the four-fold super-resolution processing effects of different models on the UIEB dataset;

[0044] Figure 6 is an ablation comparison chart of different modules for the denoising task on the UIEB dataset;

[0045] Figure 7 is an ablation comparison chart of different modules for the two-fold super-resolution task on the UIEB dataset;

[0046] Figure 8 is an ablation comparison chart of different modules for the four-fold super-resolution on the UIEB dataset;

[0047] Figure 9 is a loss curve for training models containing only shallow feature extraction modules, deep feature extraction modules, image reconstruction modules, and pre-training-fine-tuning strategies with different strategies for the denoising task;

[0048] Figure 10are loss curves when training models containing only shallow feature extraction modules, deep feature extraction modules, image reconstruction modules, and pre-training-fine-tuning strategies for different strategies for the two times super-resolution task.

[0049] Figure 11 are loss curves when training models containing only shallow feature extraction modules, deep feature extraction modules, image reconstruction modules, and pre-training-fine-tuning strategies for different strategies for the four times super-resolution task. DETAILED DESCRIPTION

[0050] The application will be further described below with reference to the accompanying drawings.

[0051] The application is a pre-training-fine-tuning method and system for underwater image restoration based on semantic and diffusion dual priors, according to Figure 1 and 2 The specific steps of the method are as follows:

[0052] Step 1: Use a blind degradation strategy to construct a pre-training dataset for pre-training;

[0053] Step 1.1: Blind degradation training strategy;

[0054] Select the original images of ImageNet100 to apply blind image degradation operations to construct training data pairs. Although this simplified degradation mode cannot completely simulate the complexity of underwater imaging, it helps the model to acquire general image degradation prior knowledge. The blind degradation includes three core components: first, in terms of blur degradation, isotropic and anisotropic Gaussian kernels are considered to simulate different types of image blur effects; second, in terms of image scale transformation, by combining region size adjustment, bilinear interpolation, and bicolor adjustment methods, diversified scale degradation is achieved; finally, in terms of noise simulation, multiple noise types such as additive Gaussian noise, Poisson noise, and JPEG compression noise are integrated. A second-order degradation mechanism is introduced, i.e., double superposition of the classic "blur-noise" degradation process to construct more complex high-order degradation effects. Each degradation has randomness and uncertainty. This comprehensive degradation strategy not only improves the robustness of the model, but also enhances its generalization ability in handling complex real scenes.

[0055] A basic pre-training network containing shallow feature extraction modules, deep feature extraction modules, and image reconstruction modules is constructed, and pre-training is performed on the ImageNet100 dataset.

[0056] Step 1.2: Shallow feature extraction module;

[0057] To effectively extract low-level visual features (such as edges, textures, and other basic visual elements) from images, a shallow feature extraction module is first constructed. Considering the significant advantages of convolutional neural networks in early visual processing, their local receptive fields and weight-sharing mechanisms effectively capture the hierarchical representation of images. A single-layer convolutional network is used to achieve efficient extraction of shallow features. Specifically, the shallow feature extraction module uses learnable convolutional kernels to map from the input space to the feature space, laying the foundation for subsequent deep feature extraction. The expression is as follows:

[0058] F shallow =f shallow (I LQ )

[0059] in, This represents a low-quality image of the original ImageNet100 image after blind image degradation (H, W, and C represent the image height, width, and number of channels, respectively), f shallow F represents the shallow feature extraction module. shallow This represents the features extracted by the shallow feature module.

[0060] Step 1.3: Deep Feature Extraction Module;

[0061] The deep feature extraction module consists of stacked residual blocks. Shallow features F shallow Through the deep feature extraction module f deep Get F deep f deep Contains N residual blocks f residual-i (i = 1, 2... N), specifically represented as follows:

[0062] F deep =f deep (F shallow )

[0063] F deep =Conv(f residual-N (f residual-(N-1) (...(f residual-1 (F shallow )))))+F shallow

[0064] Each residual block consists of a single layer of MambaVision, and each MambaVision layer contains two layers of MambaVisionMixer and one multi-layer self-attention layer.

[0065] Step 1.4: Image Reconstruction Module;

[0066] The image reconstruction module adopts a parallel double-branch architecture design, respectively for image denoising and image super-resolution two types of reconstruction tasks. The image denoising sub-task is committed to removing the scattering noise, suspended matter interference and other degradation phenomena in the underwater image, and outputting a high-quality reconstruction result with the same spatial resolution as the input image. The super-resolution task not only needs to process image noise, but also needs to reconstruct a low-resolution image into a high-resolution clear image. Both tasks adopt a residual learning strategy to learn the residual representation of image degradation through a feature extraction network, but there are significant differences in the reconstruction path: the denoising task directly realizes image repair through residual mapping and jump connection, and the super-resolution task needs an additional upsampling module to realize the resolution improvement.

[0067] The denoising task maps the features processed by the deep feature extraction module through a convolution layer, and adds it to the original input image to complete the final image reconstruction process, specifically:

[0068] I HQ_dn =I LQ +Conv(F deep )

[0069] Where I HQ_dn represents the reconstructed high-quality denoised image. Conv represents a single-layer convolution operation with a 3x3 convolution kernel. This reconstruction mechanism based on residual learning has the following advantages: by directly introducing the original input image I LQ into the reconstruction process, the main structural information of the image is preserved; the convolution layer Conv focuses on learning the degradation component of the image, and this residual mapping mechanism enables the network to more effectively capture and remove various degradation phenomena in underwater images.

[0070] On the basis of shallow feature extraction and deep feature extraction, the super-resolution task adopts an image reconstruction strategy for real scene super-resolution—nearest neighbor interpolation method to realize, for the two times super-resolution image reconstruction task:

[0071] F1=LeakyReLU(Conv(F deep ))

[0072] F2=LeakyReLU(Conv(Nearest(F1)))

[0073] I HQ_2SR =Conv(LeakyReLU(Conv(F2)))

[0074] On the basis of two times super-resolution image reconstruction, the four times super-resolution image reconstruction task:

[0075] F3=LeakyReLU(Conv(Nearest(F2)))

[0076] I HQ_4SR = Conv(LeakyReLU(Conv(F3)))

[0077] where F1, F2 and F3 represent intermediate features; I HQ_2SR , I HQ_4SR are 2x and 4x super-resolution high-quality images, respectively; Nearest represents nearest neighbor interpolation; Conv represents a single layer of convolution operation with a 3x3 convolution kernel; LeakyReLU represents a rectified linear unit with leakage. When processing low-resolution images in underwater scenes due to factors such as light attenuation, water scattering, etc., the nearest neighbor interpolation first provides fast size enlargement, and the subsequent convolution operation focuses on repairing and refining image details, such as fish scale texture, coral structure, and other fine features. This progressive processing mechanism can more accurately control feature expression during the reconstruction process compared to the one-time upsampling method. With the progressive upsampling strategy, when reconstructing a 4x super-resolution image, it is not upsampling at once, which can better control the reconstruction quality.

[0078] Step 2: Use the adapter-fine-tuning method to optimize the deep feature extraction module on the underwater image dataset UIEB;

[0079] Step 2.1: Fine-tune the adapter

[0080] An adapter is attached to each residual block in the deep feature extraction module obtained in step 1, constructing an efficient fine-tuning mechanism. The fine-tuning stage updates the adapter parameters based on the underwater image dataset to output high-quality underwater restored images. This design fully considers the balance between knowledge transfer and task adaptation: on the one hand, the adapter enhances the learning ability of the pre-trained model for specific underwater image degradation patterns by introducing additional parameter space; on the other hand, since the adapter is embedded in the original network structure in a superimposed manner, the integrity of the residual connection is maintained, effectively preserving the general image degradation prior knowledge obtained in the pre-training stage. This modular transfer learning strategy not only speeds up the adaptation process to the underwater image restoration task, but also maintains the understanding ability of general image degradation patterns.

[0081] The adapter is:

[0082] F out = F in + α(Conv(GELU(Conv(LayerNorm(F in )))))

[0083] where F in and F outInput and output features of the adapter in the deep feature extraction module, respectively. LayerNorm is a normalization technique that standardizes the feature dimension of each sample. Conv represents a depthwise separable convolution, which only calculates with one dedicated convolution kernel for each input channel. This design significantly reduces the number of parameters and computational complexity while maintaining the ability to effectively model spatial information. GELU is an activation function. Alpha is a learnable parameter with an initial value of 1e -3 .

[0084] Step 2.2: Multi-stage fine-tuning based on the progressive unfreezing strategy, and different parameter unfreezing schemes are adopted for different tasks.

[0085] Denoising task: Three-stage fine-tuning strategy is adopted: the first stage (first 20 epochs) focuses on the optimization of the adapter module, and other parameters are frozen to maintain pre-training knowledge; the second stage (20 to 100 epochs) unfreezes the image reconstruction module and the projection layer in the deep feature extraction module to increase the adaptability of the model; the final stage (100 to 700 epochs) unfreezes all modules for overall optimization.

[0086] Super-resolution task: Two-stage strategy is adopted: the first stage (first 100 epochs) optimizes the adapter module, image reconstruction module, and projection layer in the deep feature extraction module simultaneously; the final stage (100 to 700 epochs) unfreezes all parameters for overall fine-tuning. Each stage uses a decreasing learning rate strategy to ensure stable convergence of the model.

[0087] The progressive unfreezing strategy has the following advantages: first, by gradually unfreezing parameters, it effectively prevents catastrophic forgetting of pre-training knowledge; second, the differentiated unfreezing scheme fully considers the characteristics of different tasks, improving the efficiency of transfer learning; finally, the decreasing learning rate design ensures that the model can obtain stable and effective optimization at each stage.

[0088] Step 3: Semantic prior module is introduced for denoising and super-resolution tasks.

[0089] Step 3.1: Training of the semantic prior module.

[0090] The semantic prior module is inserted after the adapter of the deep feature extraction module optimized in step 2, and its core component is the SPT Unit, which remains trainable.

[0091] Firstly, the SAM model is used to generate the corresponding semantic segmentation mask for each image in the underwater dataset. Subsequently, the original low-quality underwater image and its corresponding semantic segmentation mask are jointly input into the model. In the semantic prior module, the SPT Unit generates a new feature map by fusing the following three key components: the depth features extracted by the step 2 optimized deep feature extraction module, the semantic segmentation mask corresponding to the image, and the SAM prior representation.

[0092] M = SAM (I LQ_water )

[0093] wherein, SAM represents the semantic segmentation mask generated by the SAM model; I represents the underwater low-quality image; C m represents the number of masks, which is a fixed value. When the number of semantic segmentation masks generated by the SAM model is less than C m , the semantic segmentation masks should be padded, and when the number of semantic segmentation masks generated by the SAM model is greater than C m , the semantic segmentation masks should be truncated.

[0094] P0 = Conv (ReLU (Conv (Concat (I LQ_water , M)))

[0095] wherein, Concat represents the connection operation in the channel dimension, Conv represents the convolution operation, and ReLU is an activation function. P0 represents the SAM prior representation in the initial stage.

[0096] Taking the working process of the i-th (i > 0) SPT Unit in the semantic prior module as an example:

[0097] P i+1 = Conv (ReLU (Conv (P i ))) + P i

[0098] F i ' = Conv (ReLU (Conv (Concat (F i , M)))

[0099] F i spt = Conv (ReLU (Conv (F i '))) * Conv (ReLU (Conv (P i ))) + F i '

[0100] F i n = F i + aFi spt

[0101] where P i represents the SAM prior representation needed by the i-th SPTUnit, P i+1 represents the SAM prior representation generated by the i-th SPTUnit for use by the i+1-th SPTUnit. F i represents the feature map output by the i-th residual block of the deep feature extraction module. F i represents the intermediate representation generated by the i-th SPTUnit. F i spt represents the feature map generated by the i-th SPTUnit. Linear computation is performed on the feature map generated by the i-th SPTUnit and the feature map generated by the i-th residual block to obtain the feature map F i n . α represents a constant. Finally, the semantic prior module uses the obtained feature map to reconstruct a high-quality denoised underwater image result I HQ_wdn . At the same time, the generated 2x and 4x super-resolution underwater images I HQ1_wsr2 , I HQ1_wsr4 are used as inputs for step 4.

[0102] Step 4: Further introduce a diffusion prior module for the super-resolution task;

[0103] Step 4.1: Diffusion prior module;

[0104] The diffusion prior module uses the prior knowledge of the pre-trained diffusion model to optimize the starting point selection and sampling strategy of the diffusion process, significantly improving the computational efficiency while ensuring the generation quality. The pre-trained diffusion model G θ is composed of a trainable VAE encoder E θ , a fine-tuned diffusion network ε θ , and a fixed VAE decoder D θ , where θ represents the model parameters. To maintain the original generation ability of Stable Diffusion, trainable LoRA layers are introduced in the encoder E θ and the diffusion network ε θ , which are fine-tuned using underwater images. To make the output space of the diffusion network as constant as possible, the parameters of the decoder D θ are fixed.

[0105] where Stable Diffusion ε θ is an image generation model under text condition. By extracting the text embedding (denoted as c y ) from the given text description y, the predicted noise can be represented as:

[0106]

[0107] where t ∈ {1, 2, 3,..., T} denotes the diffusion time step, T represents the total number of time steps, z t represents the diffusion state at time t.

[0108] Modify the above process, that is, the latent space transformation F θ is expressed as a text-conditioned image denoising process:

[0109]

[0110] where α T and β T respectively represent the scaling factor at the Tth diffusion time step. Only at the Tth diffusion time step, the latent representation z L of the input image is denoised in a single step, and no additional noise is introduced. The output after denoising is expected to have better realism than z L . Use the text prompt word extractor DAPE to process different resolution input images I HQ1_wsr2 , I HQ1_wsr4 , respectively, to obtain text embeddings c y1 = DAPE(I HQ1_wsr2 ), c y2 = DAPE(I HQ1_wsr4 ). Therefore, it can be expressed as:

[0111]

[0112] where I HQ2_wsr2 , I HQ2_wsr4 represent the output of the model 2 times, 4 times, the final super-resolution high-quality underwater image.

[0113] Embodiment 1:

[0114] Comparative experiment on UIEB dataset

[0115] According to Figure 3For the denoising performance of different models on the UIEB dataset, the method proposed in the present application performs better in terms of naturalness of color restoration. Other methods excessively change the original color tone, making the image appear unrealistic brownish yellow tone, such as the performance of UWGAN on the third row of images. OSMOSIS has a large difference from the reference label in multiple scenes (such as the first and fourth rows), and DiffBIR and OSEDiff methods excessively retain the color cast of the original image in some scenes. The method in this paper performs well in maintaining the details of the target object. Compared with UWGAN, it not only improves the image clarity, but also avoids the noise and artifacts caused by excessive sharpening. Especially in the surface texture of marine organisms in the third row and the details of the seabed sand and stone in the sixth row, the restoration effect of the method in this paper is more natural. In the cave scene in the fourth row, the method in this paper shows good lighting balance ability. It not only improves the visibility of the dark area inside the cave, but also maintains the natural transition of the overall lighting, avoiding the problem of overexposure of highlights or loss of dark details caused by UWGAN.

[0116] According to Figure 4 For the performance of different models on the UIEB dataset, the method proposed in the present application exhibits good detail preservation and restoration ability. Especially in the first and third rows of images, compared with the excessive smoothing or blurring effect produced by methods such as BSRNet and BSRGAN, the present application successfully preserves and enhances the texture details of the target area, making the image clearer and more realistic. In terms of naturalness of color restoration, the OSMOSIS method produces obvious color distortion in multiple scenes (such as the first, second, and sixth rows), presenting unnatural purple and red artifacts. In contrast, the present application not only effectively eliminates the color cast unique to underwater images, but also maintains the natural color distribution of the scene, avoiding the problem of artifacts caused by excessive correction. Overall, the present application shows good adaptability to different water environments (such as the shallow water area in the second row and the deep water area in the sixth row). Regardless of the turbidity of the water body, it can achieve relatively consistent image quality improvement, reflecting the robustness and generalization ability of the algorithm.

[0117] According to Figure 5 For the performance of different models on the UIEB dataset, the method proposed in the present application exhibits good detail preservation and restoration ability. Especially in the first and third rows of images, compared with the excessive smoothing or blurring effect produced by methods such as BSRNet and BSRGAN, the present application successfully preserves and enhances the texture details of the target area, making the image clearer and more realistic. In terms of naturalness of color restoration, the OSMOSIS method produces obvious color distortion in multiple scenes (such as the first, second, and sixth rows), presenting unnatural purple and red artifacts. In contrast, the present application not only effectively eliminates the color cast unique to underwater images, but also maintains the natural color distribution of the scene, avoiding the problem of artifacts caused by excessive correction. Overall, the present application shows good adaptability to different water environments (such as the shallow water area in the second row and the deep water area in the sixth row). Regardless of the turbidity of the water body, it can achieve relatively consistent image quality improvement, reflecting the robustness and generalization ability of the algorithm. Figure 4The more serious performance, showing unnatural purple and red artifacts. In contrast, the present application can still maintain the natural color distribution of the scene, avoid the problem of artifacts caused by excessive correction. Although the OSEDiff method generates images with more delicate textures, it still makes the overall image biased towards the same color tone.

[0118] Table 1 is a quantitative analysis of the denoising effect of different models on the UIEB dataset.

[0119] Table 1: Test set restoration results of each model (denoising task)

[0120] Method PSNR SSIM LPIPS MUSIQ NIQE UCIQE WaterNet 20.29 0.8955 0.1134 27.3936 6.6546 0.5566 UWGAN 20.30 0.9190 0.1400 30.8990 6.4902 0.5573 OSMOSIS 18.08 0.8908 0.1796 30.1106 6.3366 0.6674 DiffBIR 16.23 0.6118 0.3479 27.6505 6.4479 0.4940 Semi-UIR 18.69 0.8828 0.2327 29.8732 7.2923 0.5869 OSEDiff 17.86 0.5644 0.4162 24.9402 8.1298 0.4839 Ours 20.90 0.9269 0.1177 29.993 6.2061 0.5917

[0121] The experimental data in Table 1 show that the present application has certain effectiveness in the denoising task of the UIEB dataset. Compared with existing methods, compared with the WaterNet method, the present application has improved in PSNR, SSIM, MUSIQ, NIQE and UCIQE indicators, and is only slightly inferior in LPIPS indicator. Compared with the UWGAN method, the present application has improved in multiple evaluation indicators: PSNR, SSIM, LPIPS, NIQE, UCIQE. Compared with the OSMOSIS method, the present application has significantly improved in multiple key indicators: PSNR, SSIM, LPIPS, NIQE. Compared with the DiffBIR method, it can be found that the present application has improved in all indicators. Compared with the Semi-UIR method, it has improved in all indicators. Compared with the OSEDiff method, it has improved in all indicators. Through comprehensive analysis of the experimental results, it can be known that in the denoising task, even if the diffusion prior is not used, the present application is still close to the DiffBIR method in the statistical characteristics of natural images, and performs better than the DiffBIR, OSMOSIS and other methods in image similarity indicators.

[0122] Table 2 is a quantitative analysis of the effect of different models on the UIEB dataset on the two times super-resolution and four times super-resolution.

[0123] Table 2: Test set restoration results of each model (super-resolution task)

[0124]

[0125] The experimental data analysis of Table 2 shows that the method also has good performance in the super-resolution task of the UIEB dataset. In the two-fold super-resolution task, the method achieves an optimal value of 0.5436 in the UCIQE index. In the four-fold super-resolution task, the method achieves suboptimal values of 25.6893 and 0.5407 in the MUSIQ and UCIQE indexes, respectively. These experimental results show that the method proposed in the present application not only performs well in the two-fold super-resolution task, but also has good performance in the more challenging four-fold super-resolution task.

[0126] Embodiment 2

[0127] Ablation experiments on the UIEB dataset

[0128] According to Figure 6 , 7 and 8, the denoising effect, two-fold super-resolution effect and four-fold super-resolution effect of the ablation experiment are shown, wherein Ablation #1 represents not using the adapter and progressive unfreezing strategy, Ablation #2 represents using only the adapter strategy, Ablation #3 represents using the adapter and progressive unfreezing strategy, Ablation #4 represents introducing the semantic prior module, and Ablation #5 represents introducing the semantic prior module and the diffusion prior module. In the present application, only the two-stage underwater image restoration model is used for the denoising task. Table 3 is a quantitative analysis of the denoising effect of the ablation experiment.

[0129] Table 3 Influence of different modules on model performance (denoising task)

[0130] Module PSNR SSIM LPIPS MUSIQ NIQE UCIQE Ablation#1 20.49 0.9125 0.1385 29.8022 6.5725 0.5860 Ablation#2 20.43 0.9084 0.1373 29.4952 6.6538 0.5897 Ablation#3 20.72 0.9197 0.1263 29.4330 6.5158 0.5893 Ablation#4 20.9 0.9269 0.1177 29.993 6.5076 0.5917

[0131] From the quantitative evaluation results of Table 3, it can be observed that when the pre-training-fine-tuning framework is combined with the semantic prior module, the proposed model achieves optimal performance in multiple key indicators. This combination is superior to other configuration schemes.

[0132] Table 4 is a quantitative analysis of the super-resolution effect of the ablation experiment.

[0133] Table 4 Influence of different modules on model performance (super-resolution task)

[0134]

[0135] Analysis of the experimental data in Table 4 shows that, in the 2x super-resolution task, the integrated pre-training-fine-tuning framework, semantic prior module, and diffusion prior module achieve optimal performance across all evaluation metrics. This model significantly outperforms other configurations in the perceptual quality evaluation metric MUSIQ (21.82), the naturalness evaluation metric NIQE (7.1743), and the underwater image-specific evaluation metric UCIQE (0.5436). In the more challenging 4x super-resolution task, the same combination also demonstrates excellent performance in the three key metrics of MUSIQ, UIQM, and UCIQE, achieving scores of 25.68 and 7.8322, respectively. This result proves that the combined effect of the semantic prior module and the diffusion prior module can significantly improve the model's performance in underwater image super-resolution tasks, especially in high-magnification super-resolution scenarios where detail loss is severe, its advantages are even more pronounced.

[0136] according to Figure 9 This figure shows the loss curves for training models containing only shallow feature extraction, deep feature extraction, image reconstruction, and pre-training-fine-tuning strategies for denoising tasks using different strategies. In the figure, Ablation#1 represents no adapter and progressive unfreezing strategy, Ablation#2 represents using only the adapter strategy, and Ablation#3 represents using both the adapter and progressive unfreezing strategy. Figure 9 The analysis results show that the model using the pre-training-fine-tuning framework experienced significant fluctuations in the early training phase (20 epochs) when unfreezing the image reconstruction module and other networks. This is likely because the model needs to adjust the pre-trained weights to adapt to the new task and data at this stage. Therefore, the initial loss fluctuations are large. However, after 40 epochs, when the entire model is fully unfrozen, the model begins to converge rapidly, and the loss curve tends to stabilize. This indicates that the pre-training-fine-tuning framework can effectively accelerate the model's learning process and ultimately achieve better convergence and stability. The model without the pre-training-fine-tuning framework exhibits a relatively smooth loss trend throughout the training process, indicating that the model is relatively stable during training and does not show significant fluctuations. However, the relatively smooth loss curve may also mean that the model's learning process is relatively slow. The model using only the adapter and progressive unfreezing strategy converges slowly. This may be because the introduction of the progressive unfreezing strategy requires gradual adaptation of the parameters of each network layer during the initial training, thus slowing down the overall training process.

[0137] according to Figure 10 These are loss curves for training models containing only shallow feature extraction modules, deep feature extraction modules, image reconstruction modules, and pre-training-fine-tuning strategies for the 2x super-resolution task using different strategies. Figure 11are the loss curves when the model containing only the shallow feature extraction module, the deep feature extraction module, the image reconstruction module and the pre-training-fine-tuning strategy is trained with different strategies for the four times super-resolution task. Figure 10 From the analysis results of 11 As the pre-trained model is pre-trained under the condition of the denoising task, the model will face certain adaptation problems when it turns to the super-resolution task. Therefore, in the initial stage, the model using the pre-training-fine-tuning framework has a larger loss in the super-resolution task. This phenomenon shows that when the pre-trained model is migrated to a new task, it may need to go through an adaptation period to adjust the parameters in order to achieve the best performance on the new task. In the two times super-resolution task, the model using the pre-training-fine-tuning framework can quickly converge and maintain the loss at a relatively minimum value. This shows that although the initial loss of the model is larger, through the knowledge transfer of the pre-training-fine-tuning framework, the model can quickly adapt to the super-resolution task and effectively learn the image details, so it achieves better performance in the two times super-resolution task. In the four times super-resolution task, the model using the pre-training-fine-tuning framework finally converges to a similar result as the model using only the adapter module and the progressive unfreezing strategy. However, in combination with Table 4, although the model pre-trained in the denoising task does not perform as outstandingly in the four times super-resolution task as in the two times super-resolution task, it still achieves certain effectiveness.

[0138] The pre-training-fine-tuning method for underwater image restoration based on semantic and diffusion double prior proposed by the present application has good performance in PSNR, SSIM, MUSIQ and other indicators, and has certain effectiveness.

[0139] The above only describes the preferred embodiments of the present application and is not intended to limit the present application. For those skilled in the art, the present application can have various modifications and changes. Any modification, equivalent replacement, improvement, etc. made within the spirit and principles of the present application shall be included in the protection scope of the present application.

Claims

1. A pre-training and fine-tuning method for underwater image restoration based on semantic and diffusion dual priors, characterized in that: The specific steps are as follows: Step 1: Construct a pre-training dataset using a blind degradation strategy for pre-training; construct a basic pre-training network including a shallow feature extraction module, a deep feature extraction module, and an image reconstruction module; generate a training dataset from the obtained universal high-quality original image set through blind degradation operation, and train the basic pre-training network using the training dataset and the universal high-quality original image set; the image reconstruction module includes parallel denoising and super-resolution tasks. Step 2: Optimize the deep feature extraction module on the underwater image dataset using the adapter-fine-tuning method; The adapter-fine-tuning method embeds an adapter in the deep feature extraction module and utilizes the acquired low-quality underwater image dataset I. LQ_water The parameters of the adapter were fine-tuned by a multi-stage fine-tuning strategy of progressive unfreezing, thereby optimizing the deep feature extraction module in the basic pre-trained network after step 1 training. Step 3: Introduce semantic prior modules for denoising and super-resolution tasks; embed the semantic prior module after the adapter, and use the SAM model to analyze dataset I. LQ_water For each image in the dataset, a corresponding semantic segmentation mask M is generated. The deep features extracted by the optimized deep feature extraction module in step 2, the semantic segmentation mask M, and the SAM prior representation are fused using SPTUnit to obtain the feature map F. i n Using F i n The image reconstruction module in the basic pre-trained network yields the final high-quality denoised underwater image result. HQ_wdn and super-resolution high-quality underwater images I HQ1_wsr2 I HQ1_wsr4 ; Step 4: Further introduce a diffusion prior module for super-resolution tasks; employ a trainable VAE encoder E θ The finely tuned diffusion network ε θ and fixed VAE decoder D θ The pre-trained diffusion model G is composed of θ ; in the E θ and ε θ A trainable LoRA layer is introduced, and the model G is trained using a dataset of collected underwater images. θ Fine-tune the parameters of I; HQ1_wsr2 I HQ1_wsr4 Input into the fine-tuned model G θ In this process, the image embedding is generated through the text prompt word extractor DAPE, which drives single-step latent space denoising to generate the final super-resolution high-quality underwater image I. HQ2_wsr2 I HQ2_wsr4 .

2. The underwater image restoration pre-training-fine-tuning method based on semantic and diffusion dual priors according to claim 1, characterized in that: In step 1, the universal high-quality original image set undergoes blind degradation operations, which include three core parts: blur degradation, scale transformation, and noise simulation. A second-order degradation mechanism is introduced to double-superimpose the "blur-noise" degradation process.

3. The underwater image restoration pre-training-fine-tuning method based on semantic and diffusion dual priors according to claim 2, characterized in that: The training dataset first obtains shallow features F through a shallow feature extraction module. shallow Then, the deep feature F is obtained through a deep feature extraction module composed of stacked residual blocks. deep Then, the image reconstruction module performs a denoising task: I HQ_dn =I LQ +Conv(F deep ) Among them, I HQ_dn This represents the reconstructed, high-quality denoised image, where Conv represents a single-layer convolution operation with a 3×3 kernel. This represents the training dataset; The super-resolution task employs progressive nearest-neighbor interpolation upsampling, performed on a 2x super-resolution scale: F1=LeakyReLU(Conv(F deep )) F2=LeakyReLU(Conv(Nearest(F1))) I HQ_2SR =Conv(LeakyReLU(Conv(F2))) Based on the reconstruction of the 2x super-resolution image, a 4x super-resolution image reconstruction is performed: F3=LeakyReLU(Conv(Nearest(F2))) I HQ_4SR =Conv(LeakyReLU(Conv(F3))) Where F1, F2, and F3 represent intermediate features; I HQ_2SR I HQ_4SR These represent 2x and 4x super-resolution high-quality images, respectively; Nearest represents nearest-neighbor interpolation; LeakyReLU represents a rectified linear unit with leakage. Each residual block consists of a single layer of MambaVision.

4. The underwater image restoration pre-training-fine-tuning method based on semantic and diffusion dual priors according to claim 3, characterized in that: In step 2, an adapter is embedded after each residual block; the progressive unfreezing multi-stage fine-tuning strategy is as follows: in the denoising task, it is divided into three stages: the first stage optimizes the adapter parameters, and in this stage, the image reconstruction module and deep feature projection layer are unfrozen; in the final stage, all parameters are unfrozen; in the super-resolution task, it is divided into two stages: the first stage optimizes the adapter, reconstruction module and projection layer, and in the final stage, all parameters are unfrozen; in all stages, a decreasing learning rate strategy is adopted to ensure stable convergence of the model.

5. The underwater image restoration pre-training-fine-tuning method based on semantic and diffusion dual priors according to claim 4, characterized in that: The SPTUnit fusion in step 3 is as follows: Generate semantic segmentation mask M = SAM(I) LQ_water The initial SAM prior is represented as P0 = Conv(ReLU(Conv(Concat(I)); LQ_water ,M)))), where Concat represents the connection operation in the channel dimension, and ReLU represents the activation function; the output of the i-th SPTUnit is: F i n =F i +αF i spt , of which F i This represents the feature map output by the i-th residual block in the deep feature extraction module, where α represents a constant, and F... i spt =Conv(ReLU(Conv(F) i ')))*Conv(ReLU(Conv(P i )))+F i ', where P i F represents the SAM prior representation required for the i-th SPTUnit. i 'Represents the intermediate representation generated by the i-th SPTUnit.

6. The underwater image restoration pre-training-fine-tuning method based on semantic and diffusion dual priors according to claim 5, characterized in that: The single-step latent space denoising in step 4 is as follows: Where, α T With β T Let c represent the scaling factor at the T-th diffusion time step. y For text embedding, z L To obtain a latent representation of the input image at the T-th diffusion time step, the text cue word extractor DAPE is used to process the input images I at different resolutions. HQ1_wsr2 I HQ1_wsr4 , obtain text embedding c y1 =DAPE(I HQ1_wsr2 ), c y2 =DAPE(I HQ1_wsr4 The resulting super-resolution high-quality underwater image I HQ2_wsr2 I HQ2_wsr4 :

7. A pre-training and fine-tuning method for underwater image restoration based on semantic and diffusion dual priors as described in any one of claims 1-6, characterized in that: In steps 1 to 3, the loss function is MSE loss. In step 4, the loss function is a combination of MSE loss, LPIPS loss and Content loss.

8. A computer device / equipment / system, comprising a memory, a processor, and a computer program stored in the memory, characterized in that: The processor executes the computer program to implement the steps of the method according to any one of claims 1 to 7.

9. A computer-readable storage medium having a computer program / instructions stored thereon, characterized in that: When the computer program / instructions are executed by the processor, they implement the steps of the method according to any one of claims 1 to 7.

Images