add decode

mmpose/models/detectors/top_down.py

@@ -178,7 +178,7 @@ def forward_test(self, img, img_metas, return_heatmap=False, **kwargs):
 
         if self.with_keypoint:
             keypoint_result = self.keypoint_head.decode_keypoints(
-                img_metas, output_heatmap, [img_width, img_height])
+                img_metas, output_heatmap, img_size=[img_width, img_height])
             result.update(keypoint_result)
 
             if not return_heatmap:

mmpose/models/keypoint_heads/fc_head.py

@@ -115,6 +115,25 @@ def inference_model(self, x, flip_pairs=None):
             output_regression = output.detach().cpu().numpy()
         return output_regression
 
+    def decode(self, img_metas, output, **kwargs):
+        """Decode the output.
+
+        Args:
+            img_metas (list(dict)): Information about data augmentation
+                By default this includes:
+                - "image_file: path to the image file
+                - "center": center of the bbox
+                - "scale": scale of the bbox
+                - "rotation": rotation of the bbox
+                - "bbox_score": score of bbox
+            output (np.ndarray[N, K, H, W]): model predicted heatmaps.
+            kwargs: dict contains 'img_size'.
+        """
+        return self.decode_keypoints(
+            img_metas=img_metas,
+            output_regression=output,
+            img_size=kwargs['img_size'])
+
     def decode_keypoints(self, img_metas, output_regression, img_size):
         """Decode keypoints from output regression.
 

mmpose/models/keypoint_heads/top_down_base_head.py

@@ -36,7 +36,23 @@ def forward(self, **kwargs):
     def inference_model(self, **kwargs):
         """Inference function."""
 
-    def decode_keypoints(self, img_metas, output_heatmap, img_size):
+    def decode(self, img_metas, output, **kwargs):
+        """Decode the output.
+
+        Args:
+            img_metas (list(dict)): Information about data augmentation
+                By default this includes:
+                - "image_file: path to the image file
+                - "center": center of the bbox
+                - "scale": scale of the bbox
+                - "rotation": rotation of the bbox
+                - "bbox_score": score of bbox
+            output (np.ndarray[N, K, H, W]): model predicted heatmaps.
+        """
+        return self.decode_keypoints(
+            img_metas=img_metas, output_heatmap=output)
+
+    def decode_keypoints(self, img_metas, output_heatmap, **kwargs):
         """Decode keypoints from heatmaps.
 
         Args:
@@ -48,7 +64,6 @@ def decode_keypoints(self, img_metas, output_heatmap, img_size):
                 - "rotation": rotation of the bbox
                 - "bbox_score": score of bbox
             output_heatmap (np.ndarray[N, K, H, W]): model predicted heatmaps.
-            img_size (tuple(img_width, img_height)): model input image size.
         """
         batch_size = len(img_metas)