add docsting, ut

Author: Tau-J

mmpose/datasets/transforms/converting.py

@@ -25,6 +25,36 @@ class KeypointConverter(BaseTransform):
         num_keypoints (int): The number of keypoints in target dataset.
         mapping (list): A list containing mapping indexes. Each element has
             format (source_index, target_index)
+
+    Example:
+        >>> import numpy as np
+        >>> # case 1: 1-to-1 mapping
+        >>> # (0, 0) means target[0] = source[0]
+        >>> self = KeypointConverter(
+        >>>     num_keypoints=3,
+        >>>     mapping=[
+        >>>         (0, 0), (1, 1), (2, 2), (3, 3)
+        >>>     ])
+        >>> results = dict(
+        >>>     keypoints=np.arange(34).reshape(2, 3, 2),
+        >>>     keypoints_visible=np.arange(34).reshape(2, 3, 2) % 2)
+        >>> results = self(results)
+        >>> assert np.equal(results['keypoints'],
+        >>>                 np.arange(34).reshape(2, 3, 2)).all()
+        >>> assert np.equal(results['keypoints_visible'],
+        >>>                 np.arange(34).reshape(2, 3, 2) % 2).all()
+        >>>
+        >>> # case 2: 2-to-1 mapping
+        >>> # ((1, 2), 0) means target[0] = (source[1] + source[2]) / 2
+        >>> self = KeypointConverter(
+        >>>     num_keypoints=3,
+        >>>     mapping=[
+        >>>         ((1, 2), 0), (1, 1), (2, 2)
+        >>>     ])
+        >>> results = dict(
+        >>>     keypoints=np.arange(34).reshape(2, 3, 2),
+        >>>     keypoints_visible=np.arange(34).reshape(2, 3, 2) % 2)
+        >>> results = self(results)
     """
 
     def __init__(self, num_keypoints: int,

tests/test_datasets/test_transforms/test_converting.py

@@ -13,6 +13,7 @@ def setUp(self):
             img_shape=(240, 320), num_instances=4, with_bbox_cs=True)
 
     def test_transform(self):
+        # 1-to-1 mapping
         mapping = [(3, 0), (6, 1), (16, 2), (5, 3)]
         transform = KeypointConverter(num_keypoints=5, mapping=mapping)
         results = transform(self.data_info.copy())
@@ -34,3 +35,39 @@ def test_transform(self):
             self.assertTrue(
                 (results['keypoints_visible'][:, target_index] ==
                  self.data_info['keypoints_visible'][:, source_index]).all())
+
+        # 2-to-1 mapping
+        mapping = [((3, 5), 0), (6, 1), (16, 2), (5, 3)]
+        transform = KeypointConverter(num_keypoints=5, mapping=mapping)
+        results = transform(self.data_info.copy())
+
+        # check shape
+        self.assertEqual(results['keypoints'].shape[0],
+                         self.data_info['keypoints'].shape[0])
+        self.assertEqual(results['keypoints'].shape[1], 5)
+        self.assertEqual(results['keypoints'].shape[2], 2)
+        self.assertEqual(results['keypoints_visible'].shape[0],
+                         self.data_info['keypoints_visible'].shape[0])
+        self.assertEqual(results['keypoints_visible'].shape[1], 5)
+
+        # check value
+        for source_index, target_index in mapping:
+            if isinstance(source_index, tuple):
+                source_index, source_index2 = source_index
+                self.assertTrue(
+                    (results['keypoints'][:, target_index] == 0.5 *
+                     (self.data_info['keypoints'][:, source_index] +
+                      self.data_info['keypoints'][:, source_index2])).all())
+                self.assertTrue(
+                    (results['keypoints_visible'][:, target_index] ==
+                     self.data_info['keypoints_visible'][:, source_index] *
+                     self.data_info['keypoints_visible'][:,
+                                                         source_index2]).all())
+            else:
+                self.assertTrue(
+                    (results['keypoints'][:, target_index] ==
+                     self.data_info['keypoints'][:, source_index]).all())
+                self.assertTrue(
+                    (results['keypoints_visible'][:, target_index] ==
+                     self.data_info['keypoints_visible'][:,
+                                                         source_index]).all())