System and method for performing semantic image segmentation
Systems and techniques are provided for performing semantic image segmentation using a machine learning system (e.g., including one or more cross-attention transformer layers). For instance, a process can include generating one or more input image features for a frame of image data and generating one or more input depth features for a frame of depth data. One or more fused image features can be determined, at least in part, by fusing the one or more input depth features with the one or more input image features, using a first cross-attention transformer network. One or more segmentation masks can be generated for the frame of image data based on the one or more fused image features.
1. An apparatus for processing image data, comprising:
at least one memory; and
at least one processor coupled to the at least one memory, the at least one processor configured to:
generate, using an encoder transformer layer of an image encoder network, input image features for a frame of image data;
generate, using an encoder transformer layer of a depth encoder network, input depth features for a frame of depth data;
generate, using a first cross-attention transformer network configured to calculate cross-attention between the image encoder network and the depth encoder network, one or more cross-attention features based on using the input image features as a cross-attention target and the input depth features as a cross-attention source;
determine one or more fused image features at least in part by using the one or more cross-attention features to fuse the input depth features with the input image features; and
generate one or more segmentation masks for the frame of image data based on the one or more fused image features.
2. The apparatus of claim 1 , wherein, to determine the one or more fused image features, the at least one processor is configured to fuse the input depth features with the input image features based on the one or more cross-attention features and using a first cross-attention transformer branch of the first cross-attention transformer network.
3. The apparatus of claim 2 , wherein the at least one processor is further configured to:
determine one or more fused depth features at least in part by fusing the input image features with the input depth features using a second cross-attention transformer branch of the first cross-attention transformer network.
4. The apparatus of claim 3 , wherein:
the first cross-attention transformer branch is configured to determine the one or more fused image features based on calculating a first cross-attention using the input image features as a target and the input depth features as a source; and
the second cross-attention transformer branch is configured to determine the one or more fused depth features based on calculating a second cross-attention using the input depth features as a target and the input image features as a source.
5. The apparatus of claim 3 , wherein the at least one processor is further configured to:
fuse, using a second cross-attention transformer network, the one or more fused depth features with the one or more fused image features, based on calculating a unidirectional cross-attention using the one or more fused image features as a target and the one or more fused depth features as a source.
6. The apparatus of claim 5 , wherein:
the first cross-attention transformer network is a bidirectional cross-attention transformer network; and
the second cross-attention transformer network is a unidirectional cross-attention transformer network.
7. The apparatus of claim 5 , wherein the at least one processor is further configured to:
generate the one or more segmentation masks for the frame of image data based on an output of the second cross-attention transformer network, wherein the output of the second cross-attention transformer network includes the one or more fused depth features fused with the one or more fused image features.
8. The apparatus of claim 3 , wherein:
the input image features used by the first cross-attention transformer branch are the same as the input image features used by the second cross-attention transformer branch; and
the input depth features used by the first cross-attention transformer branch are the same as the input depth features used by the second cross-attention transformer branch.
9. The apparatus of claim 1 , wherein the first cross-attention transformer network:
receives the input image features from a transformer-based encoder of the image encoder network; and
receives the input depth features from a transformer-based encoder of the depth encoder network.
10. The apparatus of claim 9 , wherein:
the encoder transformer layer is included in the transformer-based encoder and is configured to generate the input image features based on the frame of image data; and
the encoder transformer layer is included in the transformer-based encoder and is configured to generate the input depth features based on the frame of depth data.
11. The apparatus of claim 9 , wherein:
the encoder transformer layer of the image encoder network generates the input image features based on previously fused image features; and
the encoder transformer layer of the depth encoder network generates the input depth features based on previously fused depth features.
12. The apparatus of claim 11 , wherein one or more of the previously fused image features and the previously fused depth features are generated at least in part by the first cross-attention transformer network.
13. The apparatus of claim 1 , wherein the frame of depth data includes a plurality of depth values for a plurality of pixels of the frame of image data.
14. The apparatus of claim 1 , wherein:
the frame of image data includes color image data; and
the frame of depth data includes a depth map associated with the frame of image data.
15. The apparatus of claim 1 , wherein the first cross-attention transformer network includes one or more shifted window transformer layers.
16. A method for processing image data, the method comprising:
generating, using an encoder transformer layer of an image encoder network, input image features for a frame of image data;
generating, using an encoder transformer layer of a depth encoder network, input depth features for a frame of depth data;
generating, using a first cross-attention transformer network configured to calculate cross-attention between the image encoder network and the depth encoder network, one or more cross-attention features based on using the input image features as a cross-attention target and the input depth features as a cross-attention source;
determining one or more fused image features at least in part by using the one or more cross-attention features to fuse the input depth features with the input image features; and
generating one or more segmentation masks for the frame of image data based on the one or more fused image features.
17. The method of claim 16 , wherein determining the one or more fused image features comprises:
fusing the input depth features with the input image features based on the one or more cross-attention features and using a first cross-attention transformer branch of the first cross-attention transformer network.
18. The method of claim 17 , further comprising:
determining one or more fused depth features at least in part by fusing the input image features with the input depth features using a second cross-attention transformer branch of the first cross-attention transformer network.
19. The method of claim 18 , wherein:
the first cross-attention transformer branch is configured to determine the one or more fused image features based on calculating a first cross-attention using the input image features as a target and the input depth features as a source; and
the second cross-attention transformer branch is configured to determine the one or more fused depth features based on calculating a second cross-attention using the input depth features as a target and the input image features as a source.
20. The method of claim 18 , further comprising:
fusing, using a second cross-attention transformer network, the one or more fused depth features with the one or more fused image features, based on calculating a unidirectional cross-attention using the one or more fused image features as a target and the one or more fused depth features as a source.
21. The method of claim 20 , wherein:
the first cross-attention transformer network is a bidirectional cross-attention transformer network; and
the second cross-attention transformer network is a unidirectional cross-attention transformer network.
22. The method of claim 20 , further comprising:
generating the one or more segmentation masks for the frame of image data based on an output of the second cross-attention transformer network, wherein the output of the second cross-attention transformer network includes the one or more fused depth features fused with the one or more fused image features.
23. The method of claim 18 , wherein:
the input image features used by the first cross-attention transformer branch are the same as the input image features used by the second cross-attention transformer branch; and
the input depth features used by the first cross-attention transformer branch are the same as the input depth features used by the second cross-attention transformer branch.
24. The method of claim 16 , wherein the first cross-attention transformer network:
receives the input image features from a transformer-based encoder of the image encoder network; and
receives the input depth features from a transformer-based encoder of the depth encoder network.
25. The method of claim 24 , wherein:
the encoder transformer layer is included in the transformer-based encoder and is configured to generate the input image features based on the frame of image data; and
the encoder transformer layer is included in the transformer-based encoder and is configured to generate the input depth features based on the frame of depth data.
26. The method of claim 24 , wherein:
the encoder transformer layer of the image encoder network generates the input image features based on previously fused image features; and
the encoder transformer layer of the depth encoder network generates the input depth features based on previously fused depth features.
27. The method of claim 26 , wherein one or more of the previously fused image features and the previously fused depth features are generated at least in part by the first cross-attention transformer network.
28. The method of claim 16 , wherein the frame of depth data includes a plurality of depth values for a plurality of pixels of the frame of image data.
29. The method of claim 16 , wherein:
the frame of image data includes color image data; and
the frame of depth data includes a depth map associated with the frame of image data.
30. The method of claim 16 , wherein the first cross-attention transformer network includes one or more shifted window transformer layers.