Two-stage prediction depth estimation and portrait matting method based on Stable Diffusion 2

Abstract

The invention belongs to the field of machine learning, depth estimation, portrait segmentation and diffusion models, and provides a two-stage prediction depth estimation and portrait matting method based on Stable Diffusion 2. According to the method, RGB images and corresponding truth values are compressed and fused into potential features through a VAE model, and first-stage coarse-grained prediction is completed through a U-Net model subjected to pre-training weight initialization; and second-stage fine-grained optimization is realized through a detail sharpening device and a multi-step correction strategy. According to the method, the uncertainty and single task limitation of an existing algorithm based on Stable Diffusion 2 are solved, the processing efficiency and the prediction precision are improved, the hardware consumption is reduced, and the method is suitable for weak texture and complex background scenes.

Description

Technical Field

[0001] This invention belongs to the fields of machine learning, depth estimation, human image segmentation, and diffusion models, and involves the diffusion model StableDiffusion2 and a two-stage coarse-to-refinement algorithm. Background Technology

[0002] Depth estimation and portrait matting are central to computer vision. With the development of convolutional neural networks and Transformers, these dense prediction algorithms have achieved significant success. In depth estimation, Clément Godar's Monodepth algorithm (2016) implemented end-to-end depth estimation for monocular images and has been continuously improved, leading to Monodepth2. In portrait segmentation, Olaf Ronneberger's U-Net and its improved U-Net++, with their encoder-decoder structure, have become classics. Modnet uses a convolutional structure to achieve instance-level portrait matting. With the rise of Transformers in computer vision, attention-based algorithms have become a research hotspot. Representative algorithms in depth estimation include the Depth Anything series and DepthPro, while ViT-Mat and other algorithms have emerged for portrait segmentation. As a core task in computer vision, the accuracy of the proposed state-of-the-art (SOTA) algorithms can meet the requirements of industrial applications. However, these algorithms are mostly applied to corresponding single dense prediction tasks. Some complex vision tasks require the combination of different basic tasks to achieve better results. Therefore, exploring a unified multi-task model is a task that the industry is currently paying more attention to.

[0003] Stable Diffusion2 is a diffusion model primarily composed of a variational autoencoder (VAE) and a U-Net. Its core idea is to follow a fixed probability flow during training, gradually and iteratively adding tiny Gaussian noise to the original data until it is "diffused" into a completely disordered random distribution. Subsequently, a deep neural network is trained to learn the reverse denoising path of this diffusion process. During inference, starting from the random noise, the model gradually removes noise based on the learned probability distribution, ultimately generating high-quality samples highly consistent with the original data distribution. This type of model, with its clear probabilistic generation mechanism, not only boasts high-quality generation and excellent detail reproduction but also strong interpretability. Therefore, it is widely used in numerous computer vision tasks such as image generation, image inpainting, image super-resolution, and image content understanding.

[0004] Stable Diffusion2 has been trained on billions of diverse image-text pairs, demonstrating geometric consistency across various scenes and an understanding of real-world physical laws. This capability suggests that the backbone network of Stable Diffusion2 implicitly contains prior knowledge of the world, fully expressing rich geometric and semantic information. Recent research has directly applied the diffusion paradigm to traditional dense prediction tasks (depth estimation and portrait matting) without reconsidering its applicability to dense prediction: the diffusion paradigm's optimization of results through random multi-step sampling contradicts the determinism and precise inference required by dense prediction tasks. Therefore, alternative methods are needed to apply Stable Diffusion2 to depth estimation and portrait matting. Summary of the Invention

[0005] This invention aims to provide a two-stage predictive depth estimation and portrait matting method based on Stable Diffusion2, which solves the problems of uncertainty and single-task limitation of existing visual algorithms based on Stable Diffusion2. At the same time, it adopts a two-stage optimization strategy to refine the depth estimation and portrait matting results.

[0006] The technical solution of this invention is as follows:

[0007] A two-stage prediction depth estimation and portrait matting method based on Stable Diffusion2 is proposed, with the following steps:

[0008] Step 1: Collect the deep dataset and the portrait matting dataset, and prepare the pre-trained weights for Stable Diffusion2;

[0009] Step 2: Input the RGB image and its corresponding depth and ground truth values ​​for portrait matting. Compress these into latent features using the VAE model, which will then serve as input to the subsequent U-Net model. Specifically:

[0010] The VAE model employs an encoder-decoder dual-branch neural network structure, which includes convolutional layers, batch normalization (BN) layers, activation function layers, and sampling layers. The encoder portion compresses the input image (e.g., 512×512) into low-dimensional latent spatial features (e.g., 64×64×4) by stacking multiple 3×3 convolutional layers (with a stride of 2 for downsampling), BN layers, and ReLU activation functions. The decoder portion, on the other hand, gradually restores the latent spatial features to an image with the same dimensions as the input image using transposed convolutional layers (with a stride of 2 for upsampling), BN layers, and activation functions. The VAE model obtains latent spatial features by processing RGB images, depth maps, and portrait matting, primarily through the following steps:

[0011] (1) Input preprocessing and alignment: The RGB image, the corresponding depth map, and the portrait cutout alpha image are made the same size. Since the VAE model can only process three-channel input, the depth map and alpha image are concatenated on the channel. The first channel is set as the depth, the second channel is set as the alpha image, and a channel is filled with all zeros to obtain the true values ​​of the three channels. If there is no portrait in the RGB, then the alpha image is filled with 0.

[0012] (2) The VAE model is used to compress the RGB image and the ground truth input respectively, compressing their resolution to 1 / 8 of the original, and obtaining the corresponding RGB image latent feature and ground truth latent feature.

[0013] (3) Post-processing: The latent space features of the RGB image obtained in step (2) and the true latent space features are cascaded and added together using the latent space feature formula:

[0014]

[0015] in Represents the latent spatial features of the RGB image. represents the true latent space feature, and t is the time step. It is the latent input feature defined after time step t.

[0016] Step 3: Initialize Stable Diffusion2 using its pre-trained weights, enabling it to possess prior world knowledge of the generative model.

[0017] Step 4: Based on the latent input features obtained from the cascade, input them into the U-Net model loaded with pre-trained weights for deterministic prediction, specifically:

[0018] (1) To avoid error accumulation in the multi-step denoising process, the original multi-step denoising formula of Stable Diffusion2 is simplified and a single-step deterministic prediction formula is adopted. Unlike the traditional noise addition-denoising logic of the diffusion model, this step does not introduce additional Gaussian noise into the latent input features, but directly uses the fused multimodal latent space features as the input of U-Net, reducing unnecessary noise interference.

[0019] (2) The mapping relationship between the ground truth latent space features and the RGB image latent space features is used as the source of supervisory information. A feature difference loss is constructed between the two:

[0020]

[0021] in This represents the mapping relationship between RGB values ​​and true values. This represents the mapping relationship predicted by the U-Net model. Through continuous training, U-Net is guided to learn an accurate mapping from RGB semantic features to depth and ground truth features, thereby improving the accuracy of the prediction results.

[0022] (3) After inference by U-Net, VAE decodes and outputs two branch results, namely the preliminary depth estimation prediction map and the portrait matting prediction result, thus completing the first stage of coarse-grained prediction (One-stage Prediction).

[0023] Step 5: Based on the coarse-grained prediction results from the first stage, refine the predictions by employing a multi-step correction strategy to geometrically refine high-frequency detail regions. Specifically, using the coarse-grained prediction results from the first stage and the original ground truth as input, repeat VAE encoding and feature concatenation and addition to feed the input latent spatial features into a U-Net with iterative loops for multi-step optimization. Finally, decode through VAE to output the refined results, namely the target depth map and portrait matting. The specific steps are as follows:

[0024] (1) Construct a detail sharpener module to learn a noiseless correction flow transformation from coarse-grained predicted features to high-fidelity ground truth features. This correction flow is defined between two known geometric states: coarse prediction and ground truth. The input to the detail sharpener is set as follows:

[0025]

[0026] in This represents coarse-grained prediction features (One-stage Prediction). This represents the features of the true latent space, thereby capturing detailed information from high-frequency detail features.

[0027] (2) During training, the number of inference loop steps is set to 10 to balance optimization and time efficiency, with flexible selection based on the required level of clarity during the inference phase. The mapping between coarse-grained predicted features and ground truth latent space features is used as the supervision source during training, with the specific training objective being:

[0028]

[0029] in Indicates the mapping relationship. This represents the prediction results of a multi-step model, which is continuously optimized by minimizing the loss.

[0030] (3) After the detail sharpener performs multi-step correction on the coarse prediction results of the first stage, it outputs the final high-precision depth estimation map and portrait matting results, thus completing the fine-grained prediction of the second stage.

[0031] The beneficial effects of this invention are:

[0032] (1) Based on the Stable Diffusion2 model framework, the end-to-end joint prediction of depth estimation and human face matting tasks is achieved through VAE multimodal feature fusion and two-stage prediction architecture. This avoids the step-by-step processing delay of traditional separate models, greatly improves the task processing efficiency in complex scenarios, and reduces the hardware resource consumption of multi-model deployment.

[0033] (2) Based on the two-stage coarse-to-fine reasoning strategy, the first stage adopts a single-step deterministic prediction formula, which reduces the error accumulation of multi-step denoising and ensures the geometric accuracy of coarse-grained prediction; the second stage uses a multi-step correction strategy and detail sharpener to target the geometric features of high-frequency detail areas, which significantly improves the spatial accuracy of the depth map and the fineness of the alpha map of portrait matting, especially suitable for weak texture and complex background scenes. Attached Figure Description

[0034] Figure 1 This is a schematic diagram of the U-Net structure of Stable Diffusion2.

[0035] Figure 2 This is a flowchart of the two-stage algorithm.

[0036] Figure 3 This is a visualization of the model's results. Detailed Implementation

[0037] The specific embodiments of the present invention will be further described below with reference to the accompanying drawings and technical solutions.

[0038] Figure 1 This is a schematic diagram of the Stable Diffusion2 U-Net architecture, a U-shaped structure of "encoder-decoder + skip connections," integrating convolutional and Transformer (attention) modules. The core consists of three parts: downsampling layers (encoder), intermediate layers (bottleneck layers), and upsampling layers (decoder layers). Each upsampling layer contains a combination of a convolutional module and a Transformer module (repeated multiple times), followed by a downsampling module to achieve feature dimensionality reduction. The intermediate layers contain convolutional and Transformer modules, processing the deepest features and fusing global information. Each upsampling layer also contains a combination of convolutional and Transformer modules, followed by an upsampling module. Simultaneously, a "concatenation" operation fuses features at the corresponding scale of the downsampling layers (skip connections). By employing a combination of convolutional and Transformer structures, important feature information can be extracted from both global and local perspectives.

[0039] This architecture employs a two-stage step-by-step prediction mechanism to achieve accurate generation of target results (depth maps, matting maps, etc.): The first stage is a single-step coarse prediction, where the input RGB image and ground truth are encoded into corresponding latent space features using a VAE, then concatenated and added before being input into a U-Net for feature processing, and finally decoded by the VAE to obtain a coarse result. The second stage is a multi-step refined prediction, using the coarse result from the first stage and the original ground truth as input, repeating the VAE encoding and feature concatenation and addition steps, then feeding the input latent space features into a U-Net with iterative loops for multi-step optimization, and finally decoding by the VAE to output the refined result, i.e., the target depth map and portrait matting. This two-stage architecture, through "coarse-to-fine" layered optimization, combined with the latent space encoding of VAE and the feature processing capabilities of U-Net, achieves a gradual improvement in result accuracy.

[0040] The training set for the model in this invention consists of the Hypersim and Virtual KITTI depth estimation training sets and the P3M-Matting portrait matting dataset, using a learning rate of 1 × 10⁻⁶. -4 The Adam optimizer is trained for 50 epochs with a gradually decaying learning rate. The first-stage predictor is set to a single-step formula (T=1) for coarse prediction, while the second-stage predictor is set to multi-step prediction (T=10) for refined prediction. For depth estimation tasks, operations are performed in the disparity space. in It is true depth.

[0041] The input image is cropped and scaled to the same resolution. The network input image size during inference is 512×512. The extracted network structure is as follows:

[0042]

[0043] We used the trained model to perform inference on the test set, and the final inference result is as follows: Figure 3 As shown.

[0044] Experimental results demonstrate that the proposed model performs exceptionally well in depth estimation tasks: it accurately identifies depth changes in dimly lit scenes and accurately extracts corresponding depth information for densely perforated scenes such as soccer nets. In portrait matting tasks, the model was validated through inference on a portrait dataset, showing not only good performance in dimly lit scenes but also complete matting of humanoid images in game scenes, fully demonstrating the model's accuracy and generalization ability.

Claims

1. A two-stage predictive depth estimation and portrait matting method based on Stable Diffusion2, characterized in that, The steps are as follows: Step 1: Collect the deep dataset and the portrait matting dataset, and prepare the pre-trained weights for Stable Diffusion2; Step 2: Input the RGB image and its corresponding depth and ground truth values ​​for portrait matting. Compress these into latent spatial features using a VAE model, which will then serve as input to the subsequent U-Net model; specifically: The VAE model employs an encoder-decoder dual-branch neural network structure, which includes convolutional layers, batch normalization layers, activation function layers, and sampling layers. The encoder portion compresses the input image into low-dimensional latent spatial features by stacking multiple 3×3 convolutional layers, batch normalization layers, and ReLU activation functions. The decoder portion then gradually restores the latent spatial features to an image with the same dimensions as the input image through transposed convolutional layers, batch normalization layers, and activation functions. The VAE model obtains latent spatial features by processing RGB images, depth maps, and portrait matting, primarily through the following steps: (1) Input preprocessing and alignment: unify the size of the RGB image, the corresponding depth map and the portrait cutout alpha image; connect the depth map and the alpha image on the channel, set the first channel as depth, set the second channel as alpha image, and add a channel with all zeros as fill to obtain the true value of the three channels. If there is no portrait in the RGB, then fill the alpha image with 0; (2) Use the VAE model to compress the RGB image and the ground truth input respectively, compressing their resolution to 1 / 8 of the original, and obtain the corresponding RGB image latent space features and ground truth latent space features. (3) Post-processing: The latent space features of the RGB image obtained in step (2) and the true latent space features are cascaded and added together using the latent space feature formula: ； in Represents the latent spatial features of the RGB image. This represents the true latent space feature, where t is the time step; These are the latent input features defined after time step t; Step 3: Initialize Stable Diffusion2 using its pre-trained weights, enabling it to possess prior world knowledge of the generative model. Step 4: Based on the latent input features obtained from the cascade, input them into the U-Net model loaded with pre-trained weights for deterministic prediction; Step 5: Based on the coarse-grained prediction results from the first stage, refine the predictions by employing a multi-step correction strategy to geometrically refine the high-frequency detail regions. Specifically, using the coarse-grained prediction results from the first stage and the original ground truth as input, repeat VAE encoding and feature concatenation and addition to feed the input latent spatial features into a U-Net with iterative loops for multi-step optimization. Finally, output the refined results, namely the target depth map and portrait matting, through VAE decoding.

2. The two-stage prediction depth estimation and portrait matting method based on Stable Diffusion2 according to claim 1, characterized in that, Step 4 is as follows: (1) The original multi-step denoising formula of Stable Diffusion2 is simplified and a single-step deterministic prediction formula is adopted; Gaussian noise is not introduced into the latent input features, and the fused multimodal latent space features are directly used as the input of U-Net; (2) The mapping relationship between the true latent space features and the RGB image latent space features is used as the source of supervisory information; a feature difference loss is constructed between the two: ； in This represents the mapping relationship between RGB values ​​and true values; This represents the mapping relationship predicted by the U-Net model. Through continuous training, U-Net is guided to learn an accurate mapping from RGB semantic features to depth and matting ground truth features, thereby improving the accuracy of the prediction results. (3) After inference by U-Net, it is decoded by VAE and outputs two branch results, namely the preliminary depth estimation prediction map and the human portrait matting prediction result, thus completing the first stage of coarse-grained prediction.

3. The two-stage prediction depth estimation and portrait matting method based on Stable Diffusion2 according to claim 2, characterized in that, Step 5 is detailed below: (1) Construct a detail sharpener module to learn a noiseless correction flow transformation from coarse-grained predicted features to high-fidelity ground truth features; this correction flow is defined between two known geometric states, coarse prediction and ground truth, and the input to the detail sharpener is set as follows: ； in This represents coarse-grained prediction features. Represents the characteristics of the truth latent space; Thus, detailed information of high-frequency detail features is captured. (2) During training, the number of inference loop steps is set to 10 to balance optimization and time efficiency, and the number of loop steps is flexibly selected according to the required level of clarity during the inference phase; during training, the mapping between coarse-grained predicted features and ground truth latent space features is used as the supervision source, and the specific training objective is: ； in Indicates the mapping relationship. This represents the prediction results of a multi-step model, which is continuously optimized by minimizing the loss. (3) After the detail sharpener performs multi-step correction on the coarse prediction results of the first stage, it outputs the final high-precision depth estimation map and portrait matting results, thus completing the fine-grained prediction of the second stage.

Images