Depth estimation based on feature reconstruction with adaptive masking and motion prediction
Systems and techniques are provided for generating depth information. For example, a process can include obtaining a first feature volume including visual features corresponding to each respective frame included in a first set of frames. A first query generator network can generate reconstruction features associated with a reconstructed feature volume corresponding to the first feature volume. Based on the first feature volume, a second query generator network can generate motion features associated with predicted future motion corresponding to the first feature volume. An initial depth prediction can be generated for each respective frame based on cross-attention between features of a depth prediction decoder, the reconstruction features, and the motion features. A refined depth prediction can be generated for each respective based on cross-attention between the initial depth prediction, the reconstruction features, and the motion features.
1 . An apparatus for generating depth information, the apparatus comprising:
one or more memories configured to store a first set of frames; and
one or more processors coupled to the one or more memories, the one or more processors being configured to:
obtain a first feature volume including visual features corresponding to each respective frame included in the first set of frames;
generate, based on the first feature volume and using a first query generator network, one or more reconstruction features associated with a reconstructed feature volume corresponding to the first feature volume;
generate, based on the first feature volume and using a second query generator network, one or more motion features associated with predicted future motion corresponding to the first feature volume;
generate an initial depth prediction for each respective frame included in the first set of frames, based on cross-attention between features of a depth prediction decoder, the one or more reconstruction features, and the one or more motion features; and
generate a refined depth prediction for each respective frame included in the first set of frames, based on cross-attention between the initial depth prediction, the one or more reconstruction features, and the one or more motion features.
2 . The apparatus of claim 1 , wherein the one or more processors are configured to:
generate, using a depth prediction decoder, the features of the depth prediction decoder, wherein the features of the depth prediction decoder are based on the first feature volume.
3 . The apparatus of claim 1 , wherein the one or more processors are configured to:
generate, using the first query generator network, an enhanced feature volume based on adaptive masking of the first feature volume; and
generate, using a depth prediction decoder, the features of the depth prediction decoder, wherein the features of the depth prediction decoder are based on the enhanced feature volume.
4 . The apparatus of claim 3 , wherein the enhanced feature volume is a reconstructed feature volume corresponding to an adaptively masked version of the first feature volume.
5 . The apparatus of claim 4 , wherein the one or more processors are configured to:
generate the adaptatively masked version of the first feature volume using an adaptive sampler machine learning network jointly trained with the first query generator network; and
generate the reconstructed feature volume based on reconstructing the first feature volume from the adaptatively masked version of the first feature volume using the first query generator network.
6 . The apparatus of claim 4 , wherein the one or more reconstruction features are associated with using the first query generator network to reconstruct the first feature volume from the adaptively masked version of the first feature volume.
7 . The apparatus of claim 1 , wherein the one or more motion features comprise attention maps obtained from a last layer of the first query generator network, based on processing the first feature volume using the first query generator network.
8 . The apparatus of claim 1 , wherein the one or more motion features comprise attention maps obtained from a last layer of the second query generator network from a final iteration step of a plurality of iteration steps performed by the second query generator network to predict a future motion corresponding to the first feature volume.
9 . The apparatus of claim 8 , wherein the second query generator network comprises a reconstruction-aware motion prediction network including a self-attention layer and a plurality of skip-attention module (SAM) layers, and wherein the last layer of the second query generator network comprises a last SAM layer of the plurality of SAM layers.
10 . The apparatus of claim 1 , wherein:
the first query generator network is a reconstruction-aware query generator machine learning network; and
the second query generator network is a motion-aware query generator machine learning network.
11 . The apparatus of claim 10 , wherein:
the reconstruction-aware query generator machine learning network is configured to generate reconstruction query features corresponding to scene features of a scene represented in the first set of frames; and
the motion-aware query generator machine learning network is configured to generate motion query features corresponding to future motion predictions of objects within the scene represented in the first set of frames.
12 . The apparatus of claim 1 , wherein, to generate the initial depth prediction, the one or more processors are configured to:
cross-attend the features of the depth prediction decoder using the one or more reconstruction features and the one or more motion features to generate cross-attended decoder features; and
generate, using the depth prediction decoder, the initial depth prediction based on the cross-attended decoder features.
13 . The apparatus of claim 12 , wherein, to generate the refined depth prediction, the one or more processors are configured to:
cross-attend, using a depth refinement machine learning network, information associated with the initial depth prediction using the one or more reconstruction features and the one or more motion features to generate cross-attended depth prediction features; and
generate, using the depth refinement machine learning network, the refined depth prediction based on the cross-attended depth prediction features.
14 . The apparatus of claim 1 , wherein:
the one or more reconstruction features associated with cross-attention for the initial depth prediction are the same as the one or more reconstruction features associated with cross-attention for the refined depth prediction; and
the one or more motion features associated with cross-attention for the initial depth prediction are the same as the one or more motion features associated with cross-attention for the refined depth prediction.
15 . The apparatus of claim 1 , wherein:
the first set of frames comprises a first set of consecutive frames of video data between a first time and a second time subsequent to the first time.
16 . The apparatus of claim 1 , wherein, to obtain the first feature volume, the one or more processors are configured to:
obtain a sequence of frames of video data including the first set of frames;
generate, using an encoder, respective visual features corresponding to each frame of the first set of frames; and
generate the first feature volume as a concatenation of the respective visual features corresponding to each frame of the first set of frames, wherein the respective visual features are concatenated in a channel dimension.
17 . The apparatus of claim 1 , further comprising one or more cameras configured to capture the first set of frames.
18 . A method for generating depth information, the method comprising:
obtaining a first feature volume including visual features corresponding to each respective frame included in a first set of frames;
generating, based on the first feature volume and using a first query generator network, one or more reconstruction features associated with a reconstructed feature volume corresponding to the first feature volume;
generating, based on the first feature volume and using a second query generator network, one or more motion features associated with predicted future motion corresponding to the first feature volume;
generating an initial depth prediction for each respective frame included in the first set of frames, based on cross-attention between features of a depth prediction decoder, the one or more reconstruction features, and the one or more motion features; and
generating a refined depth prediction for each respective frame included in the first set of frames, based on cross-attention between the initial depth prediction, the one or more reconstruction features, and the one or more motion features.
19 . The method of claim 18 , further comprising:
generating, using a depth prediction decoder, the features of the depth prediction decoder, wherein the features of the depth prediction decoder are based on the first feature volume.
20 . The method of claim 18 , further comprising:
generating, using the first query generator network, an enhanced feature volume based on adaptive masking of the first feature volume; and
generating, using a depth prediction decoder, the features of the depth prediction decoder, wherein the features of the depth prediction decoder are based on the enhanced feature volume.