An edge-guided IR-SDE image restoration method and system based on image-edge synchronous diffusion

Abstract

The present application belongs to the technical field of digital image processing, and particularly relates to an edge-guided IR-SDE image restoration method and system based on image-edge synchronous diffusion. The method of the present application is as follows: first, an edge map is extracted from a low-quality image, which is spliced with the low-quality image to form a multi-channel input; second, an image-edge synchronous mean recovery forward SDE is designed, the same diffusion process is applied to a high-quality image and a corresponding true value edge map, and synchronous modeling of double-domain degradation is realized; then, a unified U-Net network with an edge-guided fusion module is constructed; finally, the network is trained through a joint loss function, and high-quality restoration of a low-quality image is realized through reverse time SDE. Through double-domain synchronous diffusion and edge guidance throughout the whole process, the present application greatly enhances the restoration ability of image edge structure and detail texture, and realizes performance improvement in the rain removal image restoration task, while retaining the strong generalization and universal adaptation ability of the original method.

Description

Technical Field

[0001] This invention belongs to the field of digital image processing technology, specifically relating to an edge-guided IR-SDE image restoration method and system based on image-edge synchronous diffusion. Background Technology

[0002] Image restoration is a fundamental and core task in computer vision, aiming to restore high-fidelity, high-quality images from degraded, low-quality images. It is a key preprocessing step for downstream applications such as autonomous driving, medical image analysis, remote sensing interpretation, and security monitoring.

[0003] The IR-SDE method proposed in "Image Restoration with Mean-Reverting Stochastic Differential Equations" directly models the image degradation process through mean-reverting stochastic differential equations, diffusing high-quality images into low-quality images with Gaussian noise, and then restoring the image through inverse-time SDE. Combined with maximum likelihood training objectives, it achieves stable training results and general task adaptability. However, this method still has core shortcomings: it only models the diffusion in the image domain, lacking targeted modeling and constraints on edge information representing the core structure of the image, leading to problems such as blurred edges and loss of texture details in the restored image, and insufficient structural fidelity in strongly degraded scenes; the network only learns fractional functions in the image domain, unable to utilize edge priors to guide the inverse restoration process, and optimizes edge and smooth regions indiscriminately, making it difficult to balance overall visual effect with the accuracy of key structure restoration; the training objective only optimizes the trajectory in the image domain, lacking supervision and constraints in the edge domain, making the model prone to structural misalignment and edge ringing problems.

[0004] Existing technologies do not yet have a solution that deeply integrates image-edge dual-domain synchronous diffusion with the IR-SDE framework, and cannot guide the diffusion and recovery process throughout the entire process through edge priors, making it difficult to fundamentally solve the edge recovery defects of the original methods. Summary of the Invention

[0005] To address the problems and shortcomings of existing technologies, this invention aims to provide an edge-guided IR-SDE image restoration method and system based on image-edge synchronous diffusion.

[0006] To achieve the above objectives, the present invention adopts the following technical solution:

[0007] The first aspect of this invention provides an edge-guided IR-SDE image restoration method based on image-edge synchronous diffusion, comprising the following steps:

[0008] S1: Obtain pairs of low-quality images and high-quality images, generate ground truth edge maps from the high-quality images using the edge extraction module, and construct a training dataset based on the low-quality images, high-quality images, and ground truth edge maps;

[0009] S2: Extract edge maps from low-quality images using the edge extraction module, and then stitch and fuse the low-quality images and edge maps along the channel dimension to obtain multi-channel data;

[0010] S3: Construct a mean-recovery forward SDE equation for image-edge synchronization, apply the same forward diffusion process to the high-quality image and the corresponding ground truth edge map, synchronously generate a noisy image and noisy edges with fixed Gaussian noise through the forward SDE equation, and derive the closed-form solution of the forward process.

[0011] S4: Construct a unified U-Net score prediction network with an edge-guided fusion module. The network input consists of multi-channel data and time steps. Noisy images and noisy edges; deep interactive fusion of image features and edge features is performed. The network simultaneously predicts the temporal correlation fractional functions of the image domain and the edge domain, and outputs the noise prediction results of the image domain and the noise prediction results of the edge domain.

[0012] S5: Based on the training dataset constructed in step S1, the score prediction network in step S4 is trained using a joint loss function to obtain the trained image restoration model.

[0013] S6: First, generate a corresponding edge map from the low-quality image to be restored through the edge extraction module. Then, stitch the low-quality image to be restored with the edge map into a multi-channel input, which is then input into the trained image restoration model. Iterative numerical solutions are obtained through edge-guided reverse time SDE, and finally, a high-quality image is output.

[0014] Further, in step S3, the mean-recovery forward SDE equation for image-edge synchronization includes:

[0015] For the image domain, the forward SDE equation is expressed as: ;

[0016] For the marginal domain, the forward SDE equation is expressed as: ;

[0017] in, for Image domain diffusion state at time t, For low-quality images, for The edge domain diffusion state at time t, The edge map corresponding to the low-quality image; The mean rate of recovery parameter is time-dependent. For time-dependent stochastic volatility parameters, For the standard Wiener process, the image domain and the edge domain share exactly the same... , With Wiener process; constraints are , The variance is a fixed stationary variance.

[0018] Furthermore, in step S3, given an initial high-quality image... and initial truth margin map At any time Both the image domain and edge domain diffusion states follow a Gaussian distribution:

[0019] Image domain distribution:

[0020]

[0021]

[0022]

[0023] Marginal domain distribution:

[0024]

[0025]

[0026]

[0027] in, This is the cumulative mean recovery in integral form; For the diffusion state in the image domain at time t The probability density function; For the diffusion state in the image domain at time t The probability density function; It is a normal distribution with mean m and variance v; Let be the mean of the distribution of the image domain at time t; Let be the mean of the distribution of the marginal region at time t; It is the common variance of the image domain and the edge domain at time t; For initial high-quality images; This is the initial truth boundary map.

[0028] Further, in step S4, the unified U-Net score prediction network with edge-guided fusion module includes:

[0029] The network backbone adopts the U-Net architecture, which consists of three parts: encoder, bottleneck layer, and decoder. Skip connections are set between encoder and decoder, and redundant group normalization layer and self-attention layer are removed to improve inference efficiency.

[0030] Furthermore, in the unified U-Net score prediction network with edge-guided fusion module, multi-channel data and time steps of low-quality images and edge maps are stitched together. The current noisy image and noise edges are used as network input data, where the input data is first mapped to high-dimensional features through an initial convolutional layer;

[0031] A cross-attention fusion module is set up at each downsampling level of the encoder and each upsampling level of the decoder to perform cross-attention interaction between image features and edge features at the same level, thereby strengthening the guiding role of edge features on image features. The fused features then enter the next level.

[0032] The network output head is configured with a dual-branch output head, outputting the noise prediction results in the image domain and the edge domain respectively. The temporal correlation fractional function between the image domain and the edge domain is obtained through the mapping relationship between noise and the fractional function. The mapping relationship is as follows:

[0033]

[0034]

[0035] In the formula, , These represent the image domain and edge domain noise predicted by the network, respectively.

[0036] Further, in step S5, the expression for the joint loss function is:

[0037]

[0038] in, For image reconstruction loss, For edge reconstruction loss, The maximum likelihood loss is weighted by edge guidance. , This is the loss balance coefficient;

[0039] Among them, image reconstruction loss The expression is:

[0040]

[0041] In the formula, For high-quality images recovered by the model inverse, For true high-quality images;

[0042] The expression for the edge reconstruction loss is as follows:

[0043]

[0044] In the formula, The edge map is the inverse recovery of the model. For high-quality image ground truth edge maps;

[0045] Among them, the edge-guided weighted maximum likelihood loss The expression is:

[0046]

[0047] In the formula, For the total time step, The positive weighting coefficients are... This is the reverse SDE drift term predicted by the network. This is the optimal analytical solution for the reverse process. The edge weight mask matrix is ​​composed of ground-value edges from a high-quality image. Figure 2 The weights of edge regions are calculated to be 3 to 10 times that of non-edge regions. For Hadama accumulation.

[0048] Further, in step S6, the expression for the edge-guided reverse time SDE is:

[0049]

[0050] in, This is a reverse-time Wiener process. This is the image domain fractional function output by the network.

[0051] Furthermore, for the Gaussian image denoising task, an edge-guided denoising ordinary differential equation is constructed and solved deterministically, with the expression as follows:

[0052]

[0053] In the formula, The image to be denoised.

[0054] The second aspect of the present invention provides a system for edge-guided IR-SDE image restoration based on image-edge synchronous diffusion, characterized in that the system includes a data preprocessing module, an edge extraction and multi-channel fusion module, a dual-domain synchronous forward SDE modeling module, a score prediction network construction module, a training optimization module, and an inference restoration module;

[0055] The data preprocessing module is used to construct the training dataset.

[0056] The edge extraction and multi-channel fusion module is used to extract edges from the input low-quality image, complete the channel stitching of the low-quality image and the edge map, and output multi-channel data.

[0057] The dual-domain synchronous forward SDE modeling module is used to construct the mean-recovery forward SDE equation for image-edge synchronization, complete the modeling of the dual-domain degradation diffusion process, and output synchronized noisy image-noisy edge pairs, closed-form solution of diffusion process and real noise label.

[0058] The score prediction network construction module is used to construct a unified U-Net score prediction network with an edge-guided fusion module, realize deep fusion of image-edge features, and synchronously output the noise / score function prediction results of the image domain and the edge domain.

[0059] The training optimization module is used to construct a joint loss function, iteratively train and optimize the parameters of the score prediction network based on the triplet training dataset, and output the trained image restoration model.

[0060] The inference recovery module is used to perform edge extraction and multi-channel stitching on the low-quality image to be recovered. It solves the edge-guided reverse time SDE through the trained image recovery model and outputs the final high-quality image.

[0061] A third aspect of the present invention provides an electronic device, including a memory and a processor, wherein the memory stores a computer program, and the processor executes the computer program to implement the edge-guided IR-SDE image restoration method based on image-edge synchronous diffusion as described in the first aspect.

[0062] Compared with the prior art, the beneficial effects of the present invention are as follows:

[0063] (1) This invention proposes for the first time the IR-SDE framework of image-edge dual-domain synchronous diffusion, which applies a fully synchronous mean recovery SDE diffusion process to the image and the edge, realizes the joint degradation modeling of image content and structural information, and integrates the edge prior into the entire process of diffusion and recovery, fundamentally solving the core defect of insufficient edge structure recovery in the original IR-SDE method.

[0064] (2) The present invention designs a multi-channel fusion input and an edge-guided fusion module with cross attention, which directly splices low-quality images and corresponding edge maps into network input. Through cross attention, the edge features guide the image features throughout the process, allowing the network to perceive edge structure information at all time steps, which greatly improves the prediction accuracy of edge regions and the structural fidelity of reverse recovery.

[0065] (3) The present invention constructs a three-branch joint loss function of image reconstruction-edge reconstruction-edge weighted maximum likelihood, and simultaneously realizes multi-objective joint optimization of pixel-level supervision, edge structure supervision and reverse trajectory optimization. This not only ensures the overall visual effect of the restored image, but also specifically enhances the restoration accuracy of key edge regions, thus achieving a balance between perceptual quality and pixel fidelity.

[0066] (4) This invention fully retains the strong generalization and universal adaptability of the original IR-SDE method. It can be adapted to all categories of image restoration tasks by simply changing the training dataset. The edge extraction module, diffusion mode and loss function in the framework can be flexibly configured and can be customized according to different task scenarios. It has strong transferability and practicality.

[0067] (5) The present invention has achieved significant improvements in both quantitative indicators and visual effects: on the Rain100H and Rain100L rain removal datasets, compared with the original IR-SDE method, PSNR and SSIM index are improved, while LPIPS and FID index are significantly reduced. Attached Figure Description

[0068] Figure 1 This is a schematic diagram of the forward diffusion process in Embodiment 1 of the present invention;

[0069] Figure 2 This is a modified Unet network architecture diagram adapted for use in Embodiment 1 of the present invention;

[0070] Figure 3 This is a schematic diagram of the reverse diffusion process in Embodiment 1 of the present invention;

[0071] Figure 4 This refers to a high-quality image in Embodiment 1 of the present invention;

[0072] Figure 5 This is a low-quality rain map of the area to be cleared, obtained in Embodiment 1 of the present invention.

[0073] Figure 6 This refers to the rain-removed image output log document in Embodiment 1 of the present invention;

[0074] Figure 7 This is the restored de-rained image in Embodiment 1 of the present invention. Detailed Implementation

[0075] To make the objectives, technical solutions, and advantages of this invention clearer, the invention will be further described in detail below with reference to the accompanying drawings and embodiments. It should be understood that the specific embodiments described herein are merely illustrative and not intended to limit the invention.

[0076] Example 1:

[0077] An edge-guided IR-SDE image restoration method based on image-edge synchronous diffusion, the specific steps of which are as follows:

[0078] S1: Obtain pairs of low-quality images and high-quality images, and then combine the high-quality images (such as...) Figure 4 As shown, ground truth edge maps are generated through an edge extraction module, and a training dataset is constructed based on low-quality images, high-quality images, and ground truth edge maps. Specifically:

[0079] The Rain100H and Rain100L deraining datasets were used. Rain100H contains 1800 pairs of training images and 100 pairs of test images, while Rain100L contains 200 pairs of training images and 100 pairs of test images. Each pair of data contains a high-quality HQ image without rain and a low-quality LQ image with rain.

[0080] Images are read using OpenCV, converting BGR format to RGB format, converting uint8 data (0-255) to float32 tensors (0-1), and converting HWC format to CHW format. For each HQ image, after conversion to uint8 format, edge features are extracted using the cv2.Canny function, with high and low thresholds set to 50 and 150, respectively. The edge feature maps are normalized to the 0-1 range and converted to a 1×H×W tensor, which, along with the corresponding LQ and HQ images, forms the training samples. During training, samples are randomly cropped to 128×128 pixels, and random horizontal flipping, vertical flipping, and 90-degree rotation are performed with a 50% probability to improve the model's generalization ability. During testing, the original image size is maintained, and the edges are padded using the check_image_size function to ensure the image size is 2. 4 =Integer multiples of 16, adapting to downsampling operations for encoding and decoding; the dataset is constructed using the LQGTDataset class, and batch loading is completed using DataLoader, with training batch size batch_size=4, num_workers=4, and random shuffling enabled.

[0081] S2: Extract the corresponding edge map from the high-quality image using the edge extraction module. Then, concatenate and fuse the low-quality image with the edge map along the channel dimension to obtain multi-channel data. Simultaneously, use the ground truth edge map from the high-quality image as the training supervision label for the edge domain.

[0082] S3: Construct a mean-recovery forward SDE equation for image-edge synchronization. Apply the same forward diffusion process to the high-quality image and the corresponding ground truth edge map. Synchronously generate a noisy image and noisy edges with fixed Gaussian noise through the forward SDE equation, and derive the closed-form solution of the forward process, such as... Figure 1 shown. Specifically:

[0083] Using mean-recovery stochastic differential equations as the core, an identical forward diffusion process is applied to both the HQ image and its corresponding ground truth edge map. The HQ image is used as the initial state of the image domain, and the LQ image as the mean state of the image domain diffusion. Similarly, the HQ ground truth edge map is used as the initial state of the edge domain, and the corresponding edge map of the LQ image is used as the mean state of the edge domain diffusion. A noisy image with fixed Gaussian noise is generated synchronously through forward SDE. With noise edge Simultaneously, the closed-form solution of the forward process is derived; the core expression of the mean recovery forward SDE for image-edge synchronization is:

[0084] For the image domain, the forward SDE equation is expressed as: ;

[0085] For the marginal domain, the forward SDE equation is expressed as: ;

[0086] Image domain: For the initial high-quality HQ image, for Image domain diffusion state at time t, For low-quality LQ images;

[0087] Edge domain: For the initial HQ truth margin map, for The edge domain diffusion state at time t, The edge map corresponding to the LQ image;

[0088] Shared parameters: The mean rate of recovery parameter is time-dependent. For time-dependent stochastic volatility parameters, For the standard Wiener process, the image domain and the edge domain share exactly the same... , With the Wiener process, synchronous diffusion in two domains is achieved;

[0089] Constraints: ,in The variance is a fixed stationary variance;

[0090] Forward SDE has a closed-form solution, given the initial state. and At any time Both the image domain and edge domain diffusion states follow a Gaussian distribution:

[0091] Image domain distribution:

[0092]

[0093]

[0094]

[0095] Marginal domain distribution:

[0096]

[0097]

[0098]

[0099] in, The cumulative amount is the mean recovery in integral form. The image domain and the edge domain share the same cumulative amount and variance, ensuring the synchronicity of the diffusion process in both domains. For the diffusion state in the image domain at time t The probability density function; For the diffusion state in the image domain at time t The probability density function; It is a normal distribution with mean m and variance v; Let be the mean of the distribution of the image domain at time t; Let be the mean of the distribution of the marginal region at time t; It is the common variance of the image domain and the edge domain at time t; For initial high-quality images; This is the initial truth boundary map.

[0100] S4: Construct a unified U-Net score prediction network with an edge-guided fusion module, such as Figure 2 As shown, the network input consists of multi-channel data and time steps. The network performs deep interactive fusion of image features and edge features, and simultaneously predicts the temporal correlation fractional functions of the image domain and the edge domain, outputting the noise prediction results of the image domain and the noise prediction results of the edge domain.

[0101] In the unified U-Net score prediction network with edge-guided fusion module, multi-channel data and time steps of low-quality images and edge maps are stitched together. The current noisy image and noise edges are used as network input data, where the input data is first mapped to high-dimensional features through an initial convolutional layer;

[0102] A cross-attention fusion module is set up at each downsampling level of the encoder and each upsampling level of the decoder to perform cross-attention interaction between image features and edge features at the same level, thereby strengthening the guiding role of edge features on image features. The fused features then enter the next level.

[0103] The network output head is configured with a dual-branch output head, outputting the noise prediction results in the image domain and the edge domain respectively. The temporal correlation fractional function between the image domain and the edge domain is obtained through the mapping relationship between noise and the fractional function. The mapping relationship is as follows:

[0104]

[0105]

[0106] In the formula, , These represent the image domain and edge domain noise predicted by the network, respectively.

[0107] The specific structure of the unified U-Net score prediction network with edge-guided fusion module is as follows: the network backbone adopts the U-Net architecture, including an encoder, a bottleneck layer, and a decoder. Skip connections are set between the encoder and the decoder, and redundant group normalization layers and self-attention layers are removed to improve inference efficiency. Specifically:

[0108] The network adopts the ConditionalUNet structure, with the following specific configurations: input image channel number in_nc=3, output channel number out_nc=3, network base width nf=64, network depth=4, edge feature channel number edge_channels=1, and upsampling ratio upscale=1 (the rain removal task does not change the resolution).

[0109] Core implementation details:

[0110] (1) Input preprocessing and channel stitching: The network first calculates the residual between the noisy image and the conditional rain map. If an edge feature exists, then [x, cond, edge] is concatenated along the channel dimension; otherwise, [x, cond] is concatenated. The final number of channels input to the initial convolution is... .

[0111] (2) Initial convolution (init_conv): Use default_conv (standard convolution), with a kernel size of 7 and padding of 3, to map the 7-channel input features to an initial feature map of nf=64 channels.

[0112] (3) Time embedding module (time_mlp): SinusoidalPosEmb is used for sine position encoding with an encoding dimension of nf=64. After MLP processing: Linear(64, 256) -> GELU() -> Linear(256, 256), a 256-dimensional time embedding vector is finally output. Available for use by all ResBlocks.

[0113] (4) Encoder (Downs) Architecture: Total depth = 4 layers, with the number of channels doubled in each layer:

[0114] Layer i: Input dimension is .

[0115] It contains two ResBlocks: one using default_conv and the other using NonLinearity (SiLU), each ResBlock receiving a temporal embedding. It also generates scale and shift parameters to modulate the features.

[0116] Includes a Residual(PreNorm(LinearAttention)): applies linear attention after normalization to capture long-range dependencies.

[0117] It includes a downsample: except for the last layer, it uses 4×4 convolutions with a stride of 2 for downsampling; the last layer uses 1×1 convolutions to maintain resolution.

[0118] (5) Middle Bottleneck Architecture: Dimensions are The interaction and fusion of deep features are completed sequentially through: ResBlock -> Residual(PreNorm(LinearAttention)) -> ResBlock.

[0119] (6) Decoder (Ups) architecture: symmetrical to the encoder, with a total depth of 4 layers, and the number of channels in each layer is halved:

[0120] Layer i: Input dimension is .

[0121] Before each upsampling, the feature maps saved in the corresponding stage of the encoder are first concatenated in the channel dimension (SkipConnection).

[0122] It contains two ResBlocks and one Residual(PreNorm(LinearAttention)), with the same structure as the encoder.

[0123] It includes an Upsample: except for the first layer, it uses nearest neighbor upsampling plus 3×3 convolution; the first layer uses 1×1 convolution.

[0124] (7) Output Header (Final): The output features of the decoder are concatenated with the features x_ saved in the initial convolution. After a final ResBlock process, the 3-channel prediction noise map is output using nn.Conv2d(nf, out_nc, 3, 1, 1).

[0125] (8) Size check (check_image_size): At the start of the network forward propagation, check whether the height and width of the input image are within the specified limits. If the value is less than a multiple of the original size, the reflection mode is used to pad it to ensure size alignment during downsampling. Finally, the original size is cropped back to the original size during output.

[0126] S5: Design a joint loss function comprising three parts: image reconstruction loss, edge reconstruction loss, and edge-guided weighted maximum likelihood loss. Optimize the parameters of the score prediction network from step S4 using this joint loss function, guiding the network to prioritize learning features from edge regions while ensuring consistency between image and edge domain restoration. Based on the training dataset constructed in step S1, use an adaptive optimizer to iteratively train the score prediction network with the joint loss function as the optimization objective until the model loss converges, resulting in the trained edge-guided IR-SDE image restoration model. The expression for the joint loss function is:

[0127]

[0128] in, For image reconstruction loss, For edge reconstruction loss, The maximum likelihood loss is weighted by edge guidance. , The loss balance coefficient has the following value ranges: , It can be adjusted according to the task scenario;

[0129] The definitions of each sub-loss strictly follow the framework design:

[0130] Image reconstruction loss : Used to constrain the pixel-level consistency between the restored image and the true HQ image; the basic expression is:

[0131]

[0132] In the formula, For high-quality images recovered by the model inverse, For true high-quality HQ images;

[0133] Edge reconstruction loss : Used to constrain the consistency between the restored edge and the true HQ ground truth edge; the basic expression is:

[0134]

[0135] In the formula, The edge map is the inverse recovery of the model. HQ truth margin plot;

[0136] Edge-guided weighted maximum likelihood loss Based on the native IR-SDE maximum likelihood trajectory loss, higher weights are assigned to the prediction errors of image edge regions, guiding the model to prioritize optimizing the inverse recovery trajectory of edge regions. The expression is as follows:

[0137]

[0138] In the formula, For the total time step, The positive weighting coefficients are... This is the reverse SDE drift term predicted by the network. This is the optimal analytical solution for the reverse process. The edge weight mask matrix is ​​formed by the HQ ground truth edge. Figure 2 The weights of edge regions are calculated to be 3 to 10 times that of non-edge regions. For Hadamah accumulation;

[0139] Model training process: The training configuration in this embodiment is as follows: total number of iterations niter=700000, the optimizer is Adam, the initial learning rate lr_G=1e-4, beta1=0.9, beta2=0.99, weight decay weight_decay=0; the learning rate scheduler is MultiStepLR_Restart, the learning rate decay step size is [200000,400000,600000], and the decay coefficient gamma=0.5; the loss function is L1 matching loss, and the loss weight is weight=1.0; the beta of the EMA model is 0.995, and it is updated every 10 steps; a validation is performed every 5000 iterations, and the model weight with the highest PSNR in the validation set is saved.

[0140] Specific training steps:

[0141] (1) Initialization: Load the training and validation sets, and initialize the ConditionalUNet network with in_nc=3, out_nc=3, nf=64, depth=4, and edge_channels=1. Initialize the Adam optimizer, multi-step learning rate scheduler, and EMA model. Initialize the IRSDE model and set a random seed to ensure reproducible training.

[0142] (2) Forward propagation and noise prediction: For each batch of HQ images Using the LQ image cond, the noisy states at corresponding times are generated by randomly sampling time steps (timesteps) through the generate_random_states function of the IRSDE model. ; Noisy state The LQ image (cond), time steps (timesteps), and extracted edge features (edge) are input into ConditionalUNet. Internally, ConditionalUNet processes the time steps through sinusoidal position encoding, extracts features by progressive downsampling through the encoder, fuses high and low layer information through skip connections, and finally outputs prediction noise with the same size as the input through the decoder.

[0143] (3) Loss calculation: The score function is calculated based on the prediction noise. The reverse_sde_step_mean function of the IRSDE model is used to calculate the current noisy state. The score is used as input; simultaneously, through the HQ image The optimal step for the reverse process is calculated using the `reverse_optimum_step` function. The L1 loss function is used to calculate the matching loss between the predicted and optimal values. .

[0144] (4) Backpropagation and parameter update: The loss is backpropagated, the network parameters are updated through the Adam optimizer, the weights of the EMA model are updated every 10 steps, and the learning rate is updated according to the learning rate scheduler's strategy.

[0145] (5) Validation and model saving: Every 5000 iterations, inference is performed on the validation set to calculate the PSNR and SSIM indices between the rain-removed image and the HQ image, record the model weights corresponding to the optimal indices, and save the model checkpoint; every 5000 iterations, the complete training state is saved once, and recovery is supported after training is interrupted.

[0146] (6) Training termination: When the number of iterations reaches 700,000, training is terminated, and the final model weights and EMA model weights are saved.

[0147] S6: The low-quality image to be restored is first processed by the edge extraction module to generate a corresponding edge map. The low-quality image to be restored and the edge map are then concatenated into a multi-channel input, which is then input into the trained image restoration model. Iterative numerical solutions are obtained through edge-guided reverse time SDE, and finally, a high-quality image is output, such as... Figure 3 As shown. The expression for the edge-guided reverse time SDE is:

[0148]

[0149] in, This is a reverse-time Wiener process. The image domain fractional function output by the network incorporates edge prior information;

[0150] During the inference process, the Euler–Maruyama numerical method is used to iteratively solve the inverse time SDE. The initial state of the iteration is Gaussian noise with stationary variance added to the LQ image to be recovered, and the total iteration step size is 100 steps.

[0151] For Gaussian image denoising, a corresponding edge-guided denoising ordinary differential equation (ODE) can be constructed for deterministic solution, and the expression is as follows:

[0152]

[0153] In the formula, For the image to be denoised, the solution process does not require the introduction of random noise from the Wiener process, and image denoising can be completed in just 22 steps, greatly improving inference efficiency.

[0154] Based on the weights of the trained image restoration model, perform rain removal inference for a single image. Specific steps include:

[0155] (1) Image preprocessing: Obtain low-quality rain images of the area to be rained out, such as... Figure 5 As shown, the low-quality rain image (LQ) to be de-rained is read, and the format conversion from BGR to RGB is performed, normalizing the values ​​to [0, 1]. The `check_image_size` function of ConditionalUNet is called to check if the image size is 2. 4 =Multiples of 16; if not, use reflect mode for edge padding. Convert the image to a 1×3×H×W tensor and load it onto the GPU.

[0156] (2) Initial noisy state generation: Gaussian noise corresponding to max_sigma is added to the input low-quality rain map through the noise_state function of the IRSDE model to generate the initial state of the reverse process. .

[0157] (3) Rain removal based on ConditionalUNet inverse sampling: Set the conditional mean of the IRSDE model to the input low-quality rain map LQ. Load the trained ConditionalUNet (EMA) model and set it to evaluation mode. Use posterior sampling mode to perform inverse single-step updates step by step from t=T to t=1.

[0158] (4) Post-processing and output: When sampling reaches t=0, the final clear image tensor is obtained. The output tensor is numerically cropped to the range [0, 1], and the previous edge padding is removed to restore the original image size. The tensor is converted to uint8 format (0-255), the RGB channels are converted to BGR channels, and saved as the final de-rained image, such as... Figure 7 As shown.

[0159] Testing showed that this embodiment will output a rain-removed image output log document, such as... Figure 6 As shown, on the Rain100H test set, the PSNR reaches 30.596dB and the SSIM reaches 0.9147. Compared with the native IR-SDE method, both PSNR and SSIM are improved.

[0160] Example 2:

[0161] A system for edge-guided IR-SDE image restoration based on image-edge synchronous diffusion, the system comprising a data preprocessing module, an edge extraction and multi-channel fusion module, a dual-domain synchronous forward SDE modeling module, a score prediction network construction module, a training optimization module, and an inference restoration module;

[0162] The data preprocessing module is used to construct the training dataset.

[0163] The edge extraction and multi-channel fusion module is used to extract edges from the input low-quality image, complete the channel stitching of the low-quality image and the edge map, and output multi-channel data.

[0164] The dual-domain synchronous forward SDE modeling module is used to construct the mean-recovery forward SDE equation for image-edge synchronization, complete the modeling of the dual-domain degradation diffusion process, and output synchronized noisy image-noisy edge pairs, closed-form solution of diffusion process and real noise label.

[0165] The score prediction network construction module is used to construct a unified U-Net score prediction network with an edge-guided fusion module, realize deep fusion of image-edge features, and synchronously output the noise / score function prediction results of the image domain and the edge domain.

[0166] The training optimization module is used to construct a joint loss function, iteratively train and optimize the parameters of the score prediction network based on the triplet training dataset, and output the trained image restoration model.

[0167] The inference recovery module is used to perform edge extraction and multi-channel stitching on the low-quality image to be recovered. It solves the edge-guided reverse time SDE through the trained image recovery model and outputs the final high-quality image.

[0168] Example 3:

[0169] This embodiment provides an electronic device, including a memory and a processor. The memory stores a computer program, and when the processor executes the computer program, it implements the edge-guided IR-SDE image restoration method based on image-edge synchronous diffusion as described in Embodiment 1.

[0170] The hardware of the electronic device in this embodiment also includes a GPU, a display buffer memory, a RAMD / A converter, and a heat sink that work in conjunction with the processor; the GPU is responsible for processing the graphics display of the electronic device, providing image rendering and acceleration functions, and using its parallel computing advantages to accelerate the processing of large-scale data-intensive tasks.

[0171] Furthermore, the edge-guided IR-SDE image restoration method based on image-edge synchronous diffusion described in Embodiment 1 can be implemented as a computer software program. For example, this embodiment includes a computer program product comprising a computer program carried on a computer-readable medium, the computer program containing program code for performing the method. In such an embodiment, the computer program can be downloaded and installed from a network, and / or installed from a removable medium. When the computer program is executed by a processor, it performs the functions defined in the method of this application.

[0172] The above embodiments are only used to illustrate the technical solutions of the present invention, and are not intended to limit the scope of protection of the present invention. Those skilled in the art can modify or make equivalent substitutions to the technical solutions of the present invention based on the concept of the present invention, without departing from the essence and scope of the technical solutions of the present invention.

Claims

1. A method for edge-guided IR-SDE image restoration based on image-edge synchronous diffusion, characterized in that, Includes the following steps: S1: Obtain pairs of low-quality images and high-quality images, generate ground truth edge maps from the high-quality images using the edge extraction module, and construct a training dataset based on the low-quality images, high-quality images, and ground truth edge maps; S2: Extract edge maps from low-quality images using the edge extraction module, and then stitch and fuse the low-quality images and edge maps along the channel dimension to obtain multi-channel data; S3: Construct a mean-recovery forward SDE equation for image-edge synchronization, apply the same forward diffusion process to the high-quality image and the corresponding ground truth edge map, synchronously generate a noisy image and noisy edges with fixed Gaussian noise through the forward SDE equation, and derive the closed-form solution of the forward process. S4: Construct a unified U-Net score prediction network with an edge-guided fusion module. The network input consists of multi-channel data and time steps. Noisy images and noisy edges; deep interactive fusion of image features and edge features is performed. The network simultaneously predicts the temporal correlation fractional functions of the image domain and the edge domain, and outputs the noise prediction results of the image domain and the noise prediction results of the edge domain. S5: Based on the training dataset constructed in step S1, the score prediction network in step S4 is trained using a joint loss function to obtain the trained image restoration model. S6: First, generate a corresponding edge map from the low-quality image to be restored through the edge extraction module. Then, stitch the low-quality image to be restored with the edge map into a multi-channel input, which is then input into the trained image restoration model. Iterative numerical solutions are obtained through edge-guided reverse time SDE, and finally, a high-quality image is output.

2. The edge-guided IR-SDE image restoration method based on image-edge synchronous diffusion according to claim 1, characterized in that, In step S3, the mean recovery forward SDE equation for image-edge synchronization includes: For the image domain, the forward SDE equation is expressed as: ; For the marginal domain, the forward SDE equation is expressed as: ; in, for Image domain diffusion state at time t, For low-quality images, for The edge domain diffusion state at time t, The edge map corresponding to the low-quality image; The mean rate of recovery parameter is time-dependent. For time-dependent stochastic volatility parameters, For the standard Wiener process, the image domain and the edge domain share exactly the same... , With Wiener process; constraints are , The variance is a fixed stationary variance.

3. The edge-guided IR-SDE image restoration method based on image-edge synchronous diffusion according to claim 1, characterized in that, In step S3, given an initial high-quality image and initial truth margin map At any time Both the image domain and edge domain diffusion states follow a Gaussian distribution: Image domain distribution: Marginal domain distribution: in, This is the cumulative mean recovery in integral form; For the diffusion state in the image domain at time t The probability density function; For the diffusion state in the image domain at time t The probability density function; It is a normal distribution with mean m and variance v; Let be the mean of the distribution of the image domain at time t; Let be the mean of the distribution of the marginal region at time t; It is the common variance of the image domain and the edge domain at time t; For initial high-quality images; This is the initial truth boundary map.

4. The edge-guided IR-SDE image restoration method based on image-edge synchronous diffusion according to claim 1, characterized in that, In step S4, the unified U-Net score prediction network with edge-guided fusion module includes: The network backbone adopts the U-Net architecture, which consists of three parts: encoder, bottleneck layer, and decoder. Skip connections are set between encoder and decoder, and redundant group normalization layer and self-attention layer are removed to improve inference efficiency.

5. The edge-guided IR-SDE image restoration method based on image-edge synchronous diffusion according to claim 4, characterized in that, In the unified U-Net score prediction network with edge-guided fusion module, multi-channel data and time steps of low-quality images and edge maps are stitched together. The current noisy image and noise edges are used as network input data, where the input data is first mapped to high-dimensional features through an initial convolutional layer; A cross-attention fusion module is set up at each downsampling level of the encoder and each upsampling level of the decoder to perform cross-attention interaction between image features and edge features at the same level, thereby strengthening the guiding role of edge features on image features. The fused features then enter the next level. The network output head is configured with a dual-branch output head, outputting the noise prediction results in the image domain and the edge domain respectively. The temporal correlation fractional function between the image domain and the edge domain is obtained through the mapping relationship between noise and the fractional function. The mapping relationship is as follows: In the formula, , These represent the image domain and edge domain noise predicted by the network, respectively.

6. The edge-guided IR-SDE image restoration method based on image-edge synchronous diffusion according to claim 1, characterized in that, In step S5, the expression for the joint loss function is: in, For image reconstruction loss, For edge reconstruction loss, The maximum likelihood loss is weighted by edge guidance. , This is the loss balance coefficient; Among them, image reconstruction loss The expression is: In the formula, For high-quality images recovered by the model inverse, For true high-quality images; The expression for the edge reconstruction loss is as follows: In the formula, The edge map is the inverse recovery of the model. For high-quality image ground truth edge maps; Among them, the edge-guided weighted maximum likelihood loss The expression is: In the formula, For the total time step, The positive weighting coefficients are... This refers to the reverse SDE drift term predicted by the network. This is the optimal analytical solution for the reverse process. This is an edge weight mask matrix, obtained by binarizing the ground truth edge map of a high-quality image. The weight of the edge region is 3 to 10 times that of the non-edge region. For Hadama accumulation.

7. The edge-guided IR-SDE image restoration method based on image-edge synchronous diffusion according to claim 1, characterized in that, In step S6, the expression for the edge-guided reverse time SDE is: in, This is a reverse-time Wiener process. This is the image domain fractional function output by the network.

8. The edge-guided IR-SDE image restoration method based on image-edge synchronous diffusion according to any one of claims 1-6, characterized in that, For the task of Gaussian image denoising, an edge-guided denoising ordinary differential equation is constructed and solved deterministically, with the following expression: In the formula, The image to be denoised.

9. A system for edge-guided IR-SDE image restoration based on image-edge synchronous diffusion, characterized in that, The system includes a data preprocessing module, an edge extraction and multi-channel fusion module, a dual-domain synchronous forward SDE modeling module, a score prediction network construction module, a training optimization module, and an inference recovery module. The data preprocessing module is used to construct the training dataset. The edge extraction and multi-channel fusion module is used to extract edges from the input low-quality image, complete the channel stitching of the low-quality image and the edge map, and output multi-channel data. The dual-domain synchronous forward SDE modeling module is used to construct the mean-recovery forward SDE equation for image-edge synchronization, complete the modeling of the dual-domain degradation diffusion process, and output synchronized noisy image-noisy edge pairs, closed-form solution of diffusion process and real noise label. The score prediction network construction module is used to construct a unified U-Net score prediction network with an edge-guided fusion module, realize deep fusion of image-edge features, and synchronously output the noise / score function prediction results of the image domain and the edge domain. The training optimization module is used to construct a joint loss function, iteratively train and optimize the parameters of the score prediction network based on the triplet training dataset, and output the trained image restoration model. The inference recovery module is used to perform edge extraction and multi-channel stitching on the low-quality image to be recovered. It solves the edge-guided reverse time SDE through the trained image recovery model and outputs the final high-quality image.

10. An electronic device comprising a memory and a processor, wherein the memory stores a computer program, characterized in that, When the processor executes the computer program, it implements the edge-guided IR-SDE image restoration method based on image-edge synchronous diffusion as described in any one of claims 1-8.

Images