cat Keypoints

Summary: Add cat methods to Keypoints

Reviewed By: ppwwyyxx

Differential Revision: D30199199

fbshipit-source-id: 7d2c63810e7e05931c90b7ea5102a734b7ce8af4

Author: Zhipeng Fan

detectron2/structures/keypoints.py

@@ -91,6 +91,26 @@ def __repr__(self) -> str:
         s += "num_instances={})".format(len(self.tensor))
         return s
 
+    @staticmethod
+    def cat(keypoints_list: List["Keypoints"]) -> "Keypoints":
+        """
+        Concatenates a list of Keypoints into a single Keypoints
+
+        Arguments:
+            keypoints_list (list[Keypoints])
+
+        Returns:
+            Keypoints: the concatenated Keypoints
+        """
+        assert isinstance(keypoints_list, (list, tuple))
+        assert len(keypoints_list) > 0
+        assert all(isinstance(keypoints, Keypoints) for keypoints in keypoints_list)
+
+        cat_kpts = type(keypoints_list[0])(
+            torch.cat([kpts.tensor for kpts in keypoints_list], dim=0)
+        )
+        return cat_kpts
+
 
 # TODO make this nicer, this is a direct translation from C2 (but removing the inner loop)
 def _keypoints_to_heatmap(

tests/structures/test_keypoints.py

@@ -0,0 +1,19 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+import unittest
+import torch
+
+from detectron2.structures.keypoints import Keypoints
+
+
+class TestKeypoints(unittest.TestCase):
+    def test_cat_keypoints(self):
+        keypoints1 = Keypoints(torch.rand(2, 21, 3))
+        keypoints2 = Keypoints(torch.rand(4, 21, 3))
+
+        cat_keypoints = keypoints1.cat([keypoints1, keypoints2])
+        self.assertTrue(torch.all(cat_keypoints.tensor[:2] == keypoints1.tensor).item())
+        self.assertTrue(torch.all(cat_keypoints.tensor[2:] == keypoints2.tensor).item())
+
+
+if __name__ == "__main__":
+    unittest.main()