[Refactor] update mapping tables

mmpose/codecs/decoupled_heatmap.py

@@ -65,6 +65,15 @@ class DecoupledHeatmap(BaseKeypointCodec):
     # instance, so that it can assign varying sigmas based on their size
     auxiliary_encode_keys = {'bbox'}
 
+    label_mapping_table = dict(
+        keypoint_weights='keypoint_weights',
+        instance_coords='instance_coords',
+    )
+    field_mapping_table = dict(
+        heatmaps='heatmaps',
+        instance_heatmaps='instance_heatmaps',
+    )
+
     def __init__(
         self,
         input_size: Tuple[int, int],

mmpose/codecs/integral_regression_label.py

@@ -45,6 +45,12 @@ class IntegralRegressionLabel(BaseKeypointCodec):
     .. _`DSNT`: https://arxiv.org/abs/1801.07372
     """
 
+    label_mapping_table = dict(
+        keypoint_labels='keypoint_labels',
+        keypoint_weights='keypoint_weights',
+    )
+    field_mapping_table = dict(heatmaps='heatmaps', )
+
     def __init__(self,
                  input_size: Tuple[int, int],
                  heatmap_size: Tuple[int, int],

mmpose/codecs/megvii_heatmap.py

@@ -39,6 +39,9 @@ class MegviiHeatmap(BaseKeypointCodec):
     .. _`CPN`: https://arxiv.org/abs/1711.07319
     """
 
+    label_mapping_table = dict(keypoint_weights='keypoint_weights', )
+    field_mapping_table = dict(heatmaps='heatmaps', )
+
     def __init__(
         self,
         input_size: Tuple[int, int],

mmpose/codecs/msra_heatmap.py

@@ -47,6 +47,9 @@ class MSRAHeatmap(BaseKeypointCodec):
     .. _`Dark Pose`: https://arxiv.org/abs/1910.06278
     """
 
+    label_mapping_table = dict(keypoint_weights='keypoint_weights', )
+    field_mapping_table = dict(heatmaps='heatmaps', )
+
     def __init__(self,
                  input_size: Tuple[int, int],
                  heatmap_size: Tuple[int, int],

mmpose/codecs/regression_label.py

@@ -30,6 +30,11 @@ class RegressionLabel(BaseKeypointCodec):
 
     """
 
+    label_mapping_table = dict(
+        keypoint_labels='keypoint_labels',
+        keypoint_weights='keypoint_weights',
+    )
+
     def __init__(self, input_size: Tuple[int, int]) -> None:
         super().__init__()
 

mmpose/codecs/simcc_label.py

@@ -52,6 +52,12 @@ class SimCCLabel(BaseKeypointCodec):
     Estimation`: https://arxiv.org/abs/2107.03332
     """
 
+    label_mapping_table = dict(
+        keypoint_x_labels='keypoint_x_labels',
+        keypoint_y_labels='keypoint_y_labels',
+        keypoint_weights='keypoint_weights',
+    )
+
     def __init__(self,
                  input_size: Tuple[int, int],
                  smoothing_type: str = 'gaussian',

mmpose/codecs/spr.py

@@ -73,6 +73,13 @@ class SPR(BaseKeypointCodec):
         https://arxiv.org/abs/1908.09220
     """
 
+    field_mapping_table = dict(
+        heatmaps='heatmaps',
+        heatmap_weights='heatmap_weights',
+        displacements='displacements',
+        displacement_weights='displacement_weights',
+    )
+
     def __init__(
         self,
         input_size: Tuple[int, int],

mmpose/codecs/udp_heatmap.py

@@ -57,6 +57,9 @@ class UDPHeatmap(BaseKeypointCodec):
     Human Pose Estimation`: https://arxiv.org/abs/1911.07524
     """
 
+    label_mapping_table = dict(keypoint_weights='keypoint_weights', )
+    field_mapping_table = dict(heatmaps='heatmaps', )
+
     def __init__(self,
                  input_size: Tuple[int, int],
                  heatmap_size: Tuple[int, int],

mmpose/datasets/transforms/pose3d_transforms.py

@@ -113,7 +113,7 @@ def transform(self, results: Dict) -> dict:
             results[keypoints_key] = keypoints
             results[keypoints_visible_key] = keypoints_visible
             results[target_key] = lifting_target
-            results[lifting_target_visible] = lifting_target_visible
+            results['lifting_target_visible'] = lifting_target_visible
 
             # flip horizontal distortion coefficients
             if self.flip_camera:

mmpose/datasets/transforms/topdown_transforms.py

@@ -125,6 +125,8 @@ def transform(self, results: Dict) -> Optional[dict]:
             results['transformed_keypoints'] = transformed_keypoints
 
         results['input_size'] = (w, h)
+        results['input_center'] = center
+        results['input_scale'] = scale
 
         return results
 

mmpose/models/pose_estimators/topdown.py

@@ -147,12 +147,12 @@ def add_pred_to_datasample(self, batch_pred_instances: InstanceList,
             gt_instances = data_sample.gt_instances
 
             # convert keypoint coordinates from input space to image space
-            bbox_centers = gt_instances.bbox_centers
-            bbox_scales = gt_instances.bbox_scales
+            input_center = data_sample.metainfo['input_center']
+            input_scale = data_sample.metainfo['input_scale']
             input_size = data_sample.metainfo['input_size']
 
             pred_instances.keypoints = pred_instances.keypoints / input_size \
-                * bbox_scales + bbox_centers - 0.5 * bbox_scales
+                * input_scale + input_center - 0.5 * input_scale
             if 'keypoints_visible' not in pred_instances:
                 pred_instances.keypoints_visible = \
                     pred_instances.keypoint_scores

mmpose/testing/_utils.py

@@ -101,13 +101,23 @@ def get_packed_inputs(batch_size=2,
         image = rng.randint(0, 255, size=(3, h, w), dtype=np.uint8)
         inputs['inputs'] = torch.from_numpy(image)
 
+        # attributes
+        bboxes = _rand_bboxes(rng, num_instances, w, h)
+        bbox_centers, bbox_scales = bbox_xyxy2cs(bboxes)
+
+        keypoints = _rand_keypoints(rng, bboxes, num_keypoints)
+        keypoints_visible = np.ones((num_instances, num_keypoints),
+                                    dtype=np.float32)
+
         # meta
         img_meta = {
             'id': idx,
             'img_id': idx,
             'img_path': '<demo>.png',
             'img_shape': img_shape,
             'input_size': input_size,
+            'input_center': bbox_centers,
+            'input_scale': bbox_scales,
             'flip': False,
             'flip_direction': None,
             'flip_indices': list(range(num_keypoints))
@@ -119,12 +129,6 @@ def get_packed_inputs(batch_size=2,
         # gt_instance
         gt_instances = InstanceData()
         gt_instance_labels = InstanceData()
-        bboxes = _rand_bboxes(rng, num_instances, w, h)
-        bbox_centers, bbox_scales = bbox_xyxy2cs(bboxes)
-
-        keypoints = _rand_keypoints(rng, bboxes, num_keypoints)
-        keypoints_visible = np.ones((num_instances, num_keypoints),
-                                    dtype=np.float32)
 
         # [N, K] -> [N, num_levels, K]
         # keep the first dimension as the num_instances