Underwater image restoration pre-training-fine tuning method and system based on semantic and diffusion double prior, and medium

Abstract

The invention discloses an underwater image restoration pre-training-fine tuning method and system based on semantic and diffusion double prior, and a medium. The method specifically comprises the following steps: constructing a basic pre-training network containing a shallow / deep feature extraction and double-task reconstruction module; adopting a multi-degradation-mode blind degradation strategy to construct a pre-training data set, and training a model to extract generalization visual features; an adapter module is embedded, a model is finely adjusted on a UIEB data set through a progressive unfreezing strategy, specific underwater illumination attenuation, color cast distribution and scattering effects are learned, and cross-domain knowledge migration is achieved; a semantic prior module is embedded for denoising and super-resolution tasks, and image structure understanding is enhanced by fusing a semantic mask generated by SAM; a diffusion prior module is introduced for a super-resolution task, the advantages of a diffusion model in a complex image generation task are exerted, and fine noise modeling and detail reconstruction capabilities are provided for an underwater image restoration model. The method has the advantages of maintaining color balance, restoring real and fine image details and the like.

Description

Technical Field

[0001] This invention belongs to the field of image processing, specifically relating to a pre-training and fine-tuning method, system, and medium for underwater image restoration based on semantic and diffusion dual priors. Background Technology

[0002] In marine environmental monitoring and research, underwater images serve as a crucial data carrier for acquiring information on ecosystem dynamics, geological structural features, and the status of underwater engineering projects. Their quality directly impacts the reliability of research conclusions and the accuracy of engineering decisions. However, during actual imaging, 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, while refraction at the medium interface causes light path deflection. These factors collectively result in severe image quality degradation. Specifically, this manifests as blurred outlines of distant targets, loss of close-up details, color shift, and contrast reduction, hindering the practical application of underwater visual analysis techniques.

[0003] For underwater image restoration tasks, researchers have proposed two technical approaches: the first approach is based on optical physics theory, improving image quality by constructing a light propagation attenuation model and combining scattering coefficient estimation with energy compensation. The second approach leverages the nonlinear feature mapping capabilities of deep neural networks, training an end-to-end image restoration model using paired datasets. While the former has a theoretical foundation, it struggles to characterize complex environmental factors such as water turbulence and uneven distribution of suspended matter, and real-time estimation of key parameters like the scattering coefficient remains a significant challenge. These inherent limitations have led to a gradual shift in research focus towards deep learning methods, hoping to achieve better restoration results while maintaining computational efficiency through the powerful feature extraction and nonlinear modeling capabilities of neural networks.

[0004] Deep learning-based underwater image restoration methods are mainly divided into two categories based on the different structures of the backbone network: CNN-based methods and Transformer-based methods. The design of convolutional neural networks (CNNs) is inspired by the neural connection patterns of the visual cortex in animals' biological visual systems. In these networks, individual neurons respond to stimuli in local areas, and the receptive fields of different neurons partially overlap, achieving effective coverage of the entire visual field. A CNN consists of multiple convolutional layers, pooling layers, and a fully connected layer at the top. It processes two-dimensional image data by reducing data dimensionality and enhancing network robustness, efficiently extracting 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. This paper utilizes the UNet architecture for underwater image restoration. This architecture, through an encoder-decoder design, captures low-level features while restoring high-level details, making it suitable for handling complex optical degradation phenomena in underwater images. However, CNN-based methods generally suffer from limited local perception and insufficient modeling of long-range dependencies. In contrast, the Transformer model exhibits stronger advantages in capturing global information and long-range dependencies. Through its self-attention mechanism, the Transformer can more comprehensively understand the relationships between different regions in an image, and is particularly adept at processing global structural information. Transformer-based underwater image restoration methods can effectively address problems such as blurred distant objects and color distortion in underwater imaging. In 2023, Khan MDR, Kulkarni A, Phutke SS, 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 for underwater image restoration. This method focuses on utilizing a phase-based self-attention mechanism to extract uncontaminated image features and introduces an optimized phase attention module to recover structural information in the image. The phase attention module learns the attention features of the input image and accurately recovers image details, demonstrating its application potential in the field of underwater image restoration.However, Transformer-based methods generally suffer from high computational complexity. Additionally, some researchers have attempted to combine underwater image restoration models with diffusion models. For example, in 2024, Nathan OB, 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. This paper utilizes RGBD diffusion priors to restore underwater images. However, the introduction of diffusion models typically leads to excessively long inference times.

[0005] In summary, deep learning-based underwater image restoration methods can automatically learn and model the complex degradation process of underwater images. However, CNN-based methods generally suffer from limited local perception and insufficient modeling of long-range dependencies, while Transformer-based methods generally suffer from high computational complexity. Furthermore, deep learning methods face a significant challenge in underwater image restoration tasks: the lack of real, high-quality underwater images. Current solutions mainly rely on two pairing methods: one is pairing real, low-quality underwater images with artificially simulated high-quality images; the other is pairing simulated degraded underwater images with real, non-underwater, high-quality images. Both of these approaches have significant limitations. Summary of the Invention

[0006] The purpose of this invention is to provide a pre-training and fine-tuning method, system, and medium for underwater image restoration based on semantic and diffusion dual priors, which can better improve the quality of the restored underwater image.

[0007] The objective of this invention is achieved through the following technical solution:

[0008] A pre-training and fine-tuning method for underwater image restoration based on semantic and diffusion dual priors is proposed, with the following specific steps:

[0009] 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.

[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 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.

[0011] 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 the SPT Unit 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 ;

[0012] Step 4: Further introduce a diffusion prior module for super-resolution tasks; employ a trainable VAE encoder E θ A 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; adjust 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 high-resolution super-high-quality image I. HQ2_wsr2 , I HQ2_wsr4 .

[0013] Furthermore, the universal high-quality original image set in step 1 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.

[0014] Furthermore, the training dataset first obtains shallow features F through a shallow feature extraction module. shallowThen, 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:

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

[0016] 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;

[0017] The super-resolution task employs progressive nearest-neighbor interpolation upsampling, performed on a 2x super-resolution scale:

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

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

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

[0021] Based on the reconstruction of the 2x super-resolution image, a 4x super-resolution image reconstruction is performed:

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

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

[0024] 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.

[0025] Each residual block consists of a single layer of MambaVision.

[0026] Furthermore, 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, during which the image reconstruction module and deep feature projection layer are unfrozen, and 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.

[0027] Furthermore, the SPTUnit fusion in step 3 is as follows:

[0028] 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.

[0029] Furthermore, the single-step latent space denoising in step 4 is as follows:

[0030]

[0031] 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 :

[0032]

[0033] Furthermore, in steps 1 to 3, the loss function is 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 includes a memory, a processor, and a computer program stored in the memory, wherein the processor executes the computer program to implement the steps of a pre-training-fine-tuning method for underwater image restoration based on semantic and diffusion dual priors.

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

[0036] The beneficial effects of this invention are as follows:

[0037] (1) To address the scarcity of training data in underwater image restoration tasks, a pre-training and 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 various degradation modes of different degrees, enabling the model to adapt to the feature distribution of complex environments, thereby extracting more generalizable and robust visual feature representations. Based on this, the adapter-fine-tuning method is used to optimize the model on the UIEB underwater image dataset, allowing the model to learn degradation modes specific to the underwater environment, such as illumination attenuation, color shift distribution, and scattering effects, achieving efficient cross-domain knowledge transfer and thus more accurately restoring the color and details of underwater images. Experimental results on the UIEB underwater image dataset show that this method outperforms underwater image restoration models such as UWGAN and WaterNet in terms of PSNR, SSIM, and other metrics.

[0038] (2) To address the issues of semantic ambiguity and detail loss in underwater image restoration, a method based on both semantic and diffusion priors is proposed. Specifically, prior information from a semantic segmentation model (SAM) and a stable diffusion model is embedded into the MambaVision-based underwater image restoration pre-training-fine-tuning model. First, a semantic prior module is introduced for both denoising and super-resolution tasks to enhance the understanding of image semantic boundaries and target structures. Building upon this, a diffusion prior module is further introduced for super-resolution tasks, leveraging the advantages of the diffusion model in complex image generation tasks to provide the underwater image restoration model with refined noise modeling and detail reconstruction capabilities. By combining the semantic and diffusion prior modules, the model can capture global semantic features and reconstruct fine-grained visual details. Experiments on the UIEB underwater image dataset demonstrate the effectiveness of this method. Attached Figure Description

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

[0040] Figure 2 This is an overall structural diagram of the method of the present invention;

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

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

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

[0044] Figure 6 This is a comparison chart of ablation techniques for different modules in the denoising task on the UIEB dataset.

[0045] Figure 7 This is a comparison of ablation techniques for different modules in a 2x super-resolution task on the UIEB dataset.

[0046] Figure 8 This is a comparison of ablation of different modules at four times super-resolution on the UIEB dataset;

[0047] Figure 9 The loss curves are for training models containing only shallow feature extraction modules, deep feature extraction modules, image reconstruction modules, and pre-training-fine-tuning strategies for denoising tasks using different strategies.

[0048] Figure 10The loss curves are 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.

[0049] Figure 11 The loss curves are for training models containing only shallow feature extraction modules, deep feature extraction modules, image reconstruction modules, and pre-training-fine-tuning strategies for the four-fold super-resolution task using different strategies. Detailed Implementation

[0050] The present invention will now be further described with reference to the accompanying drawings.

[0051] This invention discloses a pre-training and fine-tuning method, system, and medium for underwater image restoration based on semantic and diffusion dual priors. Figure 1 and 2 The specific steps of the method are as follows:

[0052] Step 1: Construct a pre-training dataset using a blind degradation strategy for pre-training;

[0053] Step 1.1: Blind degeneracy training strategy;

[0054] Training data pairs are constructed by applying blind image degradation operations to original Images100 images. While this simplified degradation model cannot fully simulate the complexity of underwater imaging, it helps the model acquire general prior knowledge of image degradation. The blind degradation comprises 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, diverse scale degradation is achieved by combining methods such as region resizing, bilinear interpolation, and duocolor adjustment; finally, in terms of noise simulation, multiple noise types are integrated, including additive Gaussian noise, Poisson noise, and JPEG compression noise. A second-order degradation mechanism is also introduced, which involves a double superposition of the classic "blur-noise" degradation process to construct more complex higher-order degradation effects. Each degradation process exhibits randomness and uncertainty. This comprehensive degradation strategy not only improves the model's robustness but also enhances its generalization ability in handling complex real-world scenes.

[0055] We construct a basic pre-trained network that includes a shallow feature extraction module, a deep feature extraction module, and an image reconstruction module, and pre-train it 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 employs a parallel dual-branch architecture, addressing two reconstruction tasks: image denoising and image super-resolution. The image denoising sub-task aims to remove degradation phenomena such as scattering noise and suspended matter interference from underwater images, outputting a high-quality reconstruction result with the same spatial resolution as the input image. The super-resolution task not only needs to handle image noise but also reconstruct low-resolution images into high-resolution, clear ones. Both tasks utilize residual learning strategies, learning residual representations of image degradation through a feature extraction network. However, they differ significantly in their reconstruction paths: the denoising task directly achieves image inpainting through residual mapping and skip connections, while the super-resolution task requires an additional upsampling module to improve resolution.

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

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

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

[0070] Based on shallow and deep feature extraction, the super-resolution task employs a near-neighbor interpolation method—an image reconstruction strategy for real-world super-resolution—specifically for the task of reconstructing a 2x super-resolution image.

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

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

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

[0074] Building upon 2x super-resolution image reconstruction, the task of 4x super-resolution image reconstruction is as follows:

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

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

[0077] Among them, 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; Conv represents a single-layer convolution operation with a 3×3 kernel; and LeakyReLU represents a rectified linear unit with leakage. When processing low-resolution images in underwater scenes caused by factors such as light attenuation and water scattering, nearest-neighbor interpolation first provides rapid scaling, while subsequent convolution operations focus on restoring and refining image details, such as the texture of fish scales and the fine features of coral structures. This progressive processing mechanism, compared to a one-time upsampling method, allows for more precise control over feature representation during the reconstruction process. Employing a progressive upsampling strategy, when reconstructing a 4x super-resolution image, upsampling is not performed all at once, thus allowing for better control over reconstruction quality.

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

[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 during pre-training, constructing an efficient fine-tuning mechanism. This fine-tuning stage updates the adapter parameters based on the underwater image dataset, thereby outputting high-quality underwater reconstructed images. This design fully considers the balance between knowledge transfer and task adaptation: on the one hand, the adapter enhances the pre-trained model's ability to learn specific degradation patterns of underwater images by introducing additional parameter space; on the other hand, since the adapter is embedded into the original network structure in a superimposed manner, the integrity of the residual connections is maintained, thus effectively preserving the general image degradation prior knowledge obtained in the pre-training stage. This modular transfer learning strategy not only accelerates the adaptation process to the underwater image restoration task but also maintains the ability to understand general image degradation patterns.

[0081] The adapter is:

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

[0083] Among them, F in With F outThese represent the input and output features of the adapter in the deep feature extraction module, respectively. LayerNorm is a normalization technique that standardizes the feature dimensions of each sample. Conv stands for Depthwise Separable Convolution, characterized by each input channel being computed with only one dedicated convolution kernel. This design significantly reduces the number of parameters and computational cost while maintaining effective modeling of spatial information. GELU is an activation function. α is a learnable parameter with an initial value of 1e. -3 .

[0084] Step 2.2: Perform multi-stage fine-tuning based on a progressive unfreezing strategy, and adopt differentiated parameter unfreezing schemes for different tasks;

[0085] Denoising task: A three-stage fine-tuning strategy is adopted: The first stage (first 20 epochs) focuses on optimizing the adapter module while keeping other parameters frozen to maintain pre-training knowledge; the second stage (20 to 100 epochs) unfreezes the projection layers in the image reconstruction module and deep feature extraction module to increase the model's adaptability; the final stage (100 to 700 epochs) fully unfreezes all modules to achieve overall optimization.

[0086] Super-resolution task: A two-stage strategy is adopted: In the first stage (first 100 epochs), the projection layers in the adapter module, image reconstruction module, and deep feature extraction module are optimized simultaneously; in the final stage (100 to 700 epochs), all parameters are unfrozen for overall fine-tuning. A decreasing learning rate strategy is used in each stage to ensure stable convergence of the model.

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

[0088] Step 3: A semantic prior module was introduced for denoising and super-resolution tasks;

[0089] Step 3.1: Training the semantic prior module;

[0090] A semantic prior module is inserted after the adapter of the deep feature extraction module optimized in step 2. Its core component is the SPT Unit, and it is kept in a trainable state.

[0091] First, the SAM model is used to generate a corresponding semantic segmentation mask for each image in the underwater dataset. Then, the original low-quality underwater image and its corresponding semantic segmentation mask are 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 optimized deep feature extraction module in step 2, the semantic segmentation mask corresponding to the image, and the SAM prior representation.

[0092] M = SAM(I) LQ_water )

[0093] in, This represents the semantic segmentation mask generated by the SAM model; Indicates low-quality underwater images; C m This 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 When the number of semantic segmentation masks generated by the SAM model exceeds C, semantic segmentation masks should be filled in. m When doing so, the semantic segmentation mask should be truncated.

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

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

[0096] Taking the working process of the i-th (i>0) SPTUnit 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 +αFi spt

[0101] Among them, P i P represents the SAM prior representation required for the i-th SPTUnit. i+1 F represents the SAM prior representation generated by the i-th SPTUnit, for use by the (i+1)-th SPTUnit. i F represents the feature map output by the i-th residual block of the deep feature extraction module. i 'Represents the intermediate representation generated by the i-th SPTUnit. F i spt This represents the feature map generated by the i-th SPTUnit. The feature map F is obtained by performing a linear computation between the feature map generated by the i-th SPTUnit and the feature map generated by the i-th residual block. i n α represents a constant. Finally, the semantic prior module uses the obtained feature maps to reconstruct a high-quality denoised underwater image result I. HQ_wdn Simultaneously generated 2x and 4x super-resolution underwater images I HQ1_wsr2 , I HQ1_wsr4 This is used as the input for step 4.

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

[0103] Step 4.1: Diffusion Prior Module;

[0104] The diffusion prior module leverages the prior knowledge of the pre-trained diffusion model to significantly improve computational efficiency while maintaining generation quality by optimizing the selection of the starting point and sampling strategy in the diffusion process. The pre-trained diffusion model G... θ By trainable VAE encoder E θ A finely tuned diffusion network ε θ and fixed VAE decoder D θ The composition is defined by θ, where θ represents the model parameters. To maintain the original generative capabilities of StableDiffusion, the encoder E... θ and diffusion network ε θ A trainable LoRA layer is introduced, fine-tuned using underwater images. To keep the output space of the diffusion network as invariant as possible, the decoder D... θ The parameters are fixed.

[0105] Among them, StableDiffusionε θ This is a text-conditional image generation model. It extracts text embeddings (denoted as c) from a given text description y. y The predicted noise can be expressed as:

[0106]

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

[0108] Modify the above process, that is, transform the latent space F θ This can be described as a text-conditional image denoising process:

[0109]

[0110] Where, α T With β T Let z and represent the scaling factors at the T-th diffusion time step. The latent representation z of the input image is only at the T-th diffusion time step. L Perform single-step denoising without introducing any additional noise. The output after denoising is... It is expected to be more than z L It has a better sense of realism. The text prompt word extractor DAPE is used to process input images at different resolutions. HQ1_wsr2 , I HQ1_wsr4 , obtain text embedding c y1 =DAPE(I HQ1_wsr2 ), c y2 =DAPE(I HQ1_wsr4 Therefore, it can be expressed as:

[0111]

[0112] Among them, I HQ2_wsr2 , I HQ2_wsr4 This indicates that the model outputs 2x and 4x the final super-resolution high-quality underwater images.

[0113] Example 1:

[0114] Comparative experiments on the UIEB dataset

[0115] according to Figure 3The figures show the denoising performance of different models on the UIEB dataset. The method proposed in this invention demonstrates better performance in terms of the naturalness of color restoration. Other methods excessively alter the original hue, resulting in unrealistic brownish-yellow tones in the images, such as UWGAN's performance in the third row. OSMOSIS shows significant differences from the reference label in multiple scenes (e.g., rows 1 and 4), while DiffBIR and OSEDiff methods excessively preserve the color cast of the original image in some scenes. Our proposed method excels in preserving the details of the target object. Compared to UWGAN, it improves image sharpness while avoiding noise and artifacts that may result from over-sharpening. Particularly in the surface texture of marine life in row 3 and the details of seabed sand in row 6, our method achieves more natural restoration. In the cave scene in row 4, our method demonstrates good lighting balance. It improves the visibility of dark areas inside the cave while maintaining a natural transition in overall lighting, avoiding issues such as overexposure in highlights or loss of detail in shadows that may occur with UWGAN.

[0116] according to Figure 4 The figures show the performance of different models on the UIEB dataset for double super-resolution processing. The figures demonstrate that this invention exhibits better detail preservation and recovery capabilities. Particularly in the images of rows 1 and 3, compared to the over-smoothing or blurring effects produced by methods such as BSRNet and BSRGAN, this invention successfully preserves and enhances the texture details of the target region, making the image clearer and more realistic. Regarding the naturalness of color reproduction, the OSMOSIS method produces significant color distortion in multiple scenes (such as rows 1, 2, and 6), exhibiting unnatural purple and red artifacts. In contrast, this invention effectively eliminates the color cast characteristic of underwater images while maintaining the natural color distribution of the scene, avoiding artifacts caused by over-correction. Overall, this invention demonstrates good adaptability to different aquatic environments (such as the shallow water area in row 2 and the deep water area in row 6). Regardless of changes in water turbidity, it achieves relatively consistent image quality improvement, demonstrating the algorithm's robustness and generalization ability.

[0117] according to Figure 5 The figures show the performance of different models on the UIEB dataset at 4x super-resolution. The figures demonstrate that even under high-magnification super-resolution settings, this invention exhibits good restoration results. This invention still demonstrates good detail preservation and restoration capabilities at 4x super-resolution. Compared to the over-smoothing or blurring effects produced by methods such as BSRNet and BSRGAN, the restored images obtained by this invention are clearer and more realistic. The OSMOSIS method produces color distortion in the second row compared to... Figure 4The problem is more pronounced, exhibiting unnatural purple and red artifacts. In contrast, this invention maintains the natural color distribution of the scene, avoiding artifacts caused by overcorrection. While the OSEDiff method generates images with finer textures, it still tends to skew the overall image towards a single color tone.

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

[0119] Table 1. Restoration results of each model on the test set (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 invention demonstrates certain effectiveness in the denoising task of the UIEB dataset. Compared with existing methods, the present invention improves upon the WaterNet method in PSNR, SSIM, MUSIQ, NIQE, and UCIQE metrics, only slightly lagging behind in LPIPS. Compared with UWGAN, the present invention improves on multiple evaluation metrics: PSNR, SSIM, LPIPS, NIQE, and UCIQE. Compared with the OSMOSIS method, the present invention achieves significant improvements in multiple key metrics: PSNR, SSIM, LPIPS, and NIQE. Compared with the DiffBIR method, the present invention improves on all metrics. Compared with the Semi-UIR method, it improves on all metrics. Compared with the OSEDiff method, it improves on all metrics. The comprehensive analysis of the experimental results shows that, even without using diffusion priors, the present invention is still close to the DiffBIR method in terms of natural image statistical features in the denoising task, and performs better than DiffBIR, OSMOSIS and other methods in terms of image similarity index.

[0122] Table 2 presents a quantitative analysis of the performance of different models on the UIEB dataset for 2x and 4x super-resolution.

[0123] Table 2. Restoration results of each model on the test set (super-resolution task)

[0124]

[0125] The experimental data analysis in Table 2 shows that the proposed method also performs well in super-resolution tasks on the UIEB dataset. In the 2x super-resolution task, the method achieves an optimal value of 0.5436 on the UCIQE metric. In the 4x super-resolution task, the method achieves suboptimal values ​​of 25.6893 and 0.5407 on the MUSIQ and UCIQE metrics, respectively. These experimental results demonstrate that the proposed method not only performs well in the 2x super-resolution task but also exhibits excellent performance in the more challenging 4x super-resolution task.

[0126] Example 2:

[0127] Ablation experiments on the UIEB dataset

[0128] according to Figure 6 , 7 Figure 8 shows the denoising effect, 2x super-resolution effect, and 4x super-resolution effect of the ablation experiment. Ablation #1 represents no adapter and progressive unfreezing strategy; Ablation #2 represents only adapter strategy; Ablation #3 represents adapter and progressive unfreezing strategy; Ablation #4 represents introducing a semantic prior module; and Ablation #5 represents introducing a semantic prior module and a diffusion prior module. In this invention, only a two-stage underwater image restoration model is used for the denoising task. Table 3 shows the quantitative analysis of the denoising effect of the ablation experiment.

[0129] Table 3. Impact 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] As can be seen from the quantitative evaluation results in Table 3, when the pre-training-fine-tuning framework is combined with the semantic prior module, the proposed model achieves optimal performance on multiple key metrics. This combination outperforms other configuration schemes.

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

[0133] Table 4. Impact 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 11These are the 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 a 4x super-resolution task using different strategies. Figure 10 and 11 The analysis results show that, since the pre-trained model was trained under denoising task conditions, it faces certain adaptation problems when transitioning to super-resolution tasks. Therefore, in the initial stage, the model using the pre-training-fine-tuning framework exhibits a larger loss in super-resolution tasks. This phenomenon indicates that the pre-trained model may need an adaptation period to adjust parameters when transferring to a new task in order to achieve optimal performance. In the 2x super-resolution task, the pre-training-fine-tuning framework model converges quickly and maintains the loss at a relatively minimum value. This shows that although the initial loss is large, through knowledge transfer using the pre-training-fine-tuning framework, the model can quickly adapt to the super-resolution task and effectively learn image details, thus achieving good performance in the 2x super-resolution task. In the 4x super-resolution task, the pre-training-fine-tuning framework model eventually converges to a similar result to the model using only the adapter module and progressive unfreezing strategy. However, considering Table 4, although the pre-trained model performed less well in the 4x super-resolution task than in the 2x super-resolution task, it still achieved some success.

[0138] The underwater image restoration pre-training-fine-tuning method based on semantic and diffusion dual priors proposed in this invention has good performance in metrics such as PSNR, SSIM, and MUSIQ, and has a certain degree of effectiveness.

[0139] The above description is merely a preferred embodiment of the present invention and is not intended to limit the invention. Various modifications and variations can be made to the present invention by those skilled in the art. Any modifications, equivalent substitutions, improvements, etc., made within the spirit and principles of the present invention should be included within the scope of protection of the present invention.

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 θ A 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; adjust 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