Stereoscopic vision parallax prediction method based on energy function and attention fusion network

Abstract

The invention belongs to the technical field of image processing, and particularly relates to a stereoscopic vision parallax prediction method based on an energy function and an attention fusion network. The objective of the invention is to solve the problem that a real-time algorithm is difficult to obtain high precision in stereoscopic vision parallax prediction. The method comprises the following specific steps: acquiring a data set: training a network model by adopting two public virtual and real stereo matching data sets; the method comprises the following steps of: constructing a network model: constructing an Entry Function and Attention Interference network, and optimizing cost body construction and cost aggregation by utilizing an attention mechanism; designing a minimum loss function; training a network model: inputting the acquired data set into the network model for training; finely adjusting model parameters; and storing model parameters. According to the method, an end-to-end supervised method is adopted, a UNet-like network architecture is integrally used, cost body construction and cost body aggregation are improved, and the efficiency of the model is improved by using an attention mechanism, so that the real-time stereoscopic vision parallax prediction method based on deep learning is realized.

Description

TECHNICAL FIELD

[0001] The present application relates to a stereo vision disparity prediction method based on an energy function and an attention fusion network, and belongs to the technical field of image processing. BACKGROUND

[0002] Stereo vision disparity prediction is an important concept in the field of computer vision, which refers to finding corresponding points and calculating the distance and direction between them from images taken from two or more different perspectives. Its application fields include depth perception, three-dimensional reconstruction, augmented reality, autonomous driving, etc., providing important information for computer understanding and simulating the human visual system. The core challenge is to determine the disparity prediction relationship between corresponding pixels in different images. Due to the different camera perspectives and occlusion between objects, the same object may exhibit different appearances in different images, or even be partially invisible; in the face of textureless or weak texture areas, a large number of pixel points have similar or equal pixel values; when the camera is shooting under poor lighting conditions, the reflection causes the target to appear in a large area of white region. These common ill-conditioned regions make it very difficult to accurately find the disparity prediction points.

[0003] At present, traditional stereo vision disparity prediction methods are mainly divided into global disparity prediction, local disparity prediction and semi-local disparity prediction. Global disparity prediction is relatively accurate, but the calculation is complex, and the optimization result is easy to fall into local optimum; local disparity prediction is flexible and fast, but it assumes that the pixels within the local window have the same disparity, which is not always true in reality, resulting in poor performance in ill-conditioned areas; semi-global disparity prediction greatly reduces the algorithm complexity and improves the calculation speed while maintaining the accuracy of global disparity prediction, but its performance depends on the selection of artificial parameters. In the face of complex scenes, the performance of traditional stereo vision disparity prediction is difficult to satisfy. The stereo vision disparity prediction method based on deep learning can automatically adapt to various complex scenes by learning data features, with excellent precision performance, especially in the case of a large amount of data set.

[0004] However, most of the current stereo vision disparity prediction methods based on deep learning focus on precision performance, with slow speed, large demand for computing resources, and generalization. Therefore, a stereo vision disparity prediction method based on an energy function and an attention fusion network is proposed to solve the above problems. SUMMARY

[0005] (1) Technical problems solved

[0006] In view of the deficiencies of the prior art, the present application provides a stereo vision disparity prediction method based on an energy function and an attention fusion network, which solves the problems raised in the background art.

[0007] (ii) Technical Solution

[0008] The present application specifically adopts the following technical solutions to achieve the above-mentioned purposes:

[0009] A stereo vision disparity prediction method based on an energy function and an attention fusion network, comprising the following steps:

[0010] S1, obtaining a data set: training the network model using stereo vision disparity prediction data sets of two cases of public virtual scenes and real scenes, namely Scene Flow and KITTI data sets;

[0011] S2, building a network model: building an Energy Function and Attention Interaction network (EFAI-Net), using the trained MobileNetV2 network for feature extraction, constructing a global energy body as an energy function by adding a smoothing term to the basic cost volume, and then using an encoder-decoder structure with an attention mechanism to aggregate the cost and calculate the final disparity;

[0012] S3, designing a minimization loss function: training the model parameters using a smooth L1 loss function to enhance the robustness of the model;

[0013] S4, training the network model: inputting the virtual scene data set Scene Flow prepared in S1 into the network model built in S2 for training until the loss converges, which means that the model parameters have been trained, and the model parameters are saved;

[0014] S5, fine-tuning the model parameters: loading the model parameters saved in S4 with the real scene data set KITTI to train and fine-tune the model, obtaining more stable and reliable model parameters, and improving the generalization ability of the model;

[0015] S6, saving the model parameters: solidifying the finally determined model parameters and saving the model.

[0016] Further, the EFAI-Net network constructed in S2 mainly includes four parts: multi-scale feature extraction, energy function cost volume construction, attention fusion cost aggregation and final disparity prediction. The backbone network of the multi-scale feature extraction part is MobileNetV2 trained on the ImageNet dataset. Through this part, multi-scale features with sizes of 1 / 32, 1 / 16, 1 / 8 and 1 / 4 of the original binocular image can be obtained. The energy function cost volume includes a basic cost volume and a smooth cost volume. The similarity of the 1 / 4 size binocular image feature constitutes the basic cost volume, and then the features of the left image are spliced with each channel of the basic cost volume, and the smooth cost volume is obtained after the self-attention mechanism calculation. The basic cost volume contains the similarity information of the binocular image, and the smooth cost volume obtains the effective content information through the self-attention mechanism. Adding the two obtains an energy function cost volume with rich information. The basic framework of cost aggregation is realized by using a three-dimensional convolution-based encoder-decoder structure. Among them, the energy function cost volume is down-sampled to 1 / 32 of the original image size by using 6 consecutive three-dimensional convolutions, and then the accurate prediction disparity level probability is obtained by using 3 consecutive up-sampling blocks. Each up-sampling block is composed of an attention fusion module and a transposed convolution. The attention fusion module splices the semantic features of the multi-scale feature extraction and the low-resolution cost volume along the channel, enhances the effective geometric features through the self-attention mechanism, suppresses the invalid part, uses the semantic features to guide the up-sampling of the low-resolution cost volume, and obtains the accurate prediction disparity level probability. Finally, the winner-takes-all strategy is used to select the disparity level with the maximum probability value as the pixel disparity, reducing the calculation amount of the model.

[0017] (III) Advantages

[0018] Compared with the prior art, the present application provides a stereo vision disparity prediction method based on an energy function and an attention fusion network, which has the following advantages:

[0019] 1. Using pre-trained MobileNetV2 for multi-scale feature extraction, compared with ResNet-like used in traditional deep learning, it has smaller parameter amount, faster speed and better multi-scale feature extraction effect. Using pre-trained parameters for transfer learning effectively reduces the training time.

[0020] 2. Using the smooth volume containing content information as the smooth item in the basic cost volume to construct the energy function cost volume, so that the cost volume has both the similarity information and the content information of the binocular image features, and does not have a large amount of redundant information, to ensure the real-time performance and accuracy of the network.

[0021] 3. Use the extracted multi-scale features and geometric features in the cost aggregation to perform attention fusion to guide the cost aggregation to better learn local features and enhance the network’s performance in disparity prediction in ill-conditioned regions such as occluded regions, weak texture regions and reflective regions.

[0022] 4. Cost aggregation is performed using 3D convolutions with an encoder-decoder structure, which enables the network to capture features in the disparity dimension, reducing the number of network parameters and computational cost, and improving the network's robustness and generalization.

[0023] 5. Using a winner-takes-all strategy to estimate parallax further reduces computational costs and enables the network to have better real-time performance. Attached Figure Description

[0024] Figure 1 This is a schematic diagram of the testing process of the present invention;

[0025] Figure 2 This is a diagram of the overall network architecture of the present invention;

[0026] Figure 3 This is a diagram of the multi-scale feature extraction structure of the present invention;

[0027] Figure 4 This is a diagram of the bottleneck structure of the present invention;

[0028] Figure 5 This is a structural diagram of the energy function cost body of the present invention;

[0029] Figure 6 This is a structural diagram of the encoder-decoder of the present invention;

[0030] Figure 7 This is a structural diagram of the attention fusion module of the present invention. Detailed Implementation

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

[0032] Example

[0033] like Figure 1 As shown, a stereo vision disparity prediction method based on energy function and attention fusion network is presented. The method specifically includes the following steps:

[0034] S1, Dataset Preparation: The entire network is trained using two datasets: publicly available virtual scene datasets and real-world scene datasets, namely Scene Flow and KITTI scene datasets. The Scene Flow dataset contains 35,454 training stereo image pairs and 4,370 test stereo image pairs for traffic scenes; KITTI is a real-world driving scene dataset containing 394 training stereo image pairs and 395 test stereo image pairs for road scenes, with each image cropped to 256×512 pixels.

[0035] S2, Construct the network model. For example... Figure 2 As shown, the overall network model structure consists of four parts: multi-scale feature extraction, energy function cost volume construction, cost aggregation, and disparity prediction. The structure diagram of multi-scale feature extraction is shown below. Figure 3 As shown, it consists of a portion of MobileNetV2, deconvolution blocks, and an output block. Bottleneck blocks with 16, 24, 32, 96, and 160 channels from MobileNetV2 are used to compress the feature map to 1 / 32 of the original image size. Each bottleneck block consists of a 1×1 kernel convolution with a stride of 1 followed by a ReLU function, a 3×3 kernel DW convolution with a stride of 2 followed by a ReLU function, and a 1×1 kernel convolution with a stride of 1 followed by a Linear function. This inverse residual structure effectively reduces the number of parameters and computational cost, such as... Figure 4 As shown. The energy function cost body structure diagram is as follows. Figure 5 As shown, the cost body includes a basic cost body and a smooth cost body, which are added together to obtain the energy function cost body. The basic cost body is obtained by calculating the similarity of a quarter-sized stereo image feature map along the channels in eight groups. Simultaneously, the quarter-sized stereo image features are fused and expanded to the same size as the basic cost body, and the two are mixed to obtain a hybrid cost body. A multi-head self-attention module is added after the hybrid cost body to obtain self-attention weights. Multiplying this by the hybrid cost body yields a smooth cost body containing effective content information. The information-rich energy function cost body is obtained by adding the basic cost body containing similarity information and the smooth cost body containing effective content information. Cost aggregation is an encoder-decoder structure containing three downsampling operations and three upsampling operations, as shown in the diagram below. Figure 6 As shown. Each upsampling step includes a 3×3×3 kernel convolution with a stride of 2 followed by batch normalization, and another 3×3×3 kernel convolution with a stride of 1 followed by batch normalization. Each upsampling module contains an attention fusion module and a transposed convolution. (See diagram.) Figure 7As shown, the attention fusion module concatenates the multi-scale features obtained from multi-scale feature extraction with the low-resolution cost volume along the channel dimension to obtain a hybrid feature cost volume. After the hybrid feature cost volume, a multi-dimensional attention fusion module performs self-attention fusion from three dimensions: disparity, channel, and space. It extracts effective information by compressing the four-dimensional hybrid feature cost volume in different dimensions, and then expands it back to its original size to avoid loss of effective information. The multi-dimensional attention fusion module reduces redundant information in the hybrid feature cost volume and enhances its effective information. Simultaneously, the attention fusion module connects the upsampling part of cost aggregation with the upsampling part of feature extraction, which helps to better learn geometric features during backpropagation. Disparity calculation uses a winner-take-all strategy, selecting the disparity level with the highest probability output from cost aggregation as the pixel disparity.

[0036] S3, Design the loss function to minimize it. The optimal network model is obtained by updating the network parameters to minimize the loss function. The robust, smooth L1 loss function is chosen, and its definition is as follows:

[0037]

[0038] The final loss function is as follows:

[0039]

[0040] Where d0 represents the predicted disparity map, d gt This represents the true parallax map.

[0041] S4, Training the Network Model. The binocular images from the dataset are input into the network to generate predicted disparity maps. Training is performed on the virtual dataset SceneFlow, with 40 training iterations and 8 images per batch. The batch size can be determined based on the computer's graphics processing unit (GPU) performance; a slightly larger batch size is better if performance allows, as it enables faster and more stable network convergence. The learning rate is set to 0.002, and halved at iterations 20, 24, 28, 32, and 36 to ensure fast network fitting without overfitting.

[0042] The Adam algorithm is chosen as the network parameter optimizer. This algorithm dynamically adjusts the learning rate for different parameters by calculating the first and second moment estimates of the gradient, exhibiting good numerical stability and fast convergence speed. When the network's loss value fluctuates within a small range, the network can be considered to have converged, meaning the network training is complete.

[0043] S5, Fine-tuning the model. The model is trained and fine-tuned using the real outdoor road dataset KITTI to obtain stable and reliable model parameters, further improving the model's ability to predict disparity.

[0044] S6, Save the model. After the network is trained, all parameters in the network need to be saved. When stereo vision disparity prediction is needed for binocular images, the binocular images can be directly input into the network to obtain the final disparity image.

[0045] The implementation of convolution, activation functions, attention mechanisms, and residual structures are algorithms well-known to those skilled in the art, and the specific processes and methods can be found in relevant textbooks or technical literature.

[0046] This invention achieves real-time stereo disparity prediction with competitive accuracy by constructing an energy function cost body and an attention fusion network model. In the experimental scheme examples of this invention, the selected disparity prediction evaluation metrics are: Peak Signal-to-Noise Ratio (PSNR), Structural Similarity Index (SSIM), Floating-Point Operations (FLOPs), and Runtime. The Mobilestereonet-2D stereo disparity prediction method mentioned in the paper "Mobilestereonet: Towards lightweight deep networks for stereo matching" is selected for comparison. Mean endpoint error refers to the average pixel difference between the predicted disparity and the true disparity; a smaller value indicates more accurate prediction. Peak Signal-to-Noise Ratio is an objective indicator of image quality, evaluated by calculating the pixel difference between the original and distorted images; a higher value indicates less image distortion. Floating-point operations measure the total number of floating-point operations required by a model or algorithm during execution. The experimental results of different stereo disparity prediction methods are shown in Table 1.

[0047] Table 1 Comparison of experimental results

[0048]

[0049] As can be seen from the table, the method proposed in this invention has fewer floating-point operations, shorter running time, and higher peak signal-to-noise ratio and structural similarity index, further demonstrating that this invention has better disparity prediction results.

[0050] Finally, it should be noted that the above descriptions are merely preferred embodiments of the present invention and are not intended to limit the present invention. Although the present invention has been described in detail with reference to the foregoing embodiments, those skilled in the art can still modify the technical solutions described in the foregoing embodiments or make equivalent substitutions for some of the technical features. Any modifications, equivalent substitutions, improvements, etc., made within the spirit and principles of the present invention should be included within the protection scope of the present invention.

Claims

1. A method for predicting disparity in stereo vision based on an energy function cost volume and an attention fusion network, characterized in that: The method includes the following steps: S1, Obtain the dataset: Train the network model using publicly available stereo vision disparity prediction datasets for both virtual and real-world scenarios, namely Scene Flow and KITTI datasets. S2, Building the network model: Build the Energy Function and Attention Interaction network (EFAI-Net), use the trained MobileNetV2 network for feature extraction, construct the global energy body as the energy function by adding a smoothing term to the basic cost body, and then use the encoder-decoder structure that integrates the attention mechanism to aggregate the cost and calculate the final disparity; S3, Design the minimum loss function: Use the smooth L1 loss function to train the model parameters to enhance the robustness of the model; S4, Train the network model: Input the virtual scene dataset Scene Flow prepared in S1 into the network model built in S2 for training until the loss converges. The model parameters are considered to have been trained and the model parameters are saved. S5, Fine-tuning model parameters: Load the model parameters saved in S4 using the real-world KITTI dataset, train and fine-tune the model to obtain more stable and reliable model parameters, and improve the model's generalization ability; S6, Save model parameters: Solidify the finalized model parameters and save the model.

2. The stereo vision disparity prediction method based on energy function cost volume and attention fusion network according to claim 1, characterized in that: The Energy Function and Attention Interaction Network (EFAI-Net) constructed in S2 uses a pre-trained MobileNetV2 network as the backbone network for feature extraction. After downsampling through multiple bottlenecks, it uses three transposed convolutions for upsampling to obtain four multi-scale features of different sizes. A basic cost body is constructed using the similarity of binocular image features. Self-attention is calculated using the content and similarity information of binocular image features to obtain a smooth cost body containing effective content information. The two are added together to obtain an information-rich energy function cost body. The cost body is aggregated using the encoder-decoder structure. In the decoder step, the semantic features in the multi-scale features are fused with the encoded low-resolution cost body to obtain a hybrid feature body. The channel spatial attention of the hybrid feature body is calculated using the attention module CBAM. After transposed convolution, accurate high-resolution geometric features are obtained. The disparity level probability of each pixel is calculated using the geometric features obtained by cost aggregation. A winner-take-all strategy is used to reduce the computational cost of the network and obtain accurate disparity prediction.

Images