bottom-up (open-mmlab#5)

Co-authored-by: wangcan <[email protected]>
Spiideo · Jul 11, 2020 · 925741f · 925741f
1 parent 3a46a40
commit 925741f
Show file tree

Hide file tree

Showing 29 changed files with 1,974 additions and 268 deletions.
diff --git a/docs/benchmark.md b/docs/benchmark.md
@@ -15,7 +15,7 @@ The training speed is measure with s/iter. The lower, the better.
 ### Results on COCO val2017 with detector having human AP of 56.4 on COCO val2017 dataset
 
 We demonstrate the superiority of our MMPose framework in terms of speed and accuracy on the standard COCO keypoint detection benchmark.
-The mAP (the mean average precision) is used as the evaluation metrics.
+The mAP (the mean average precision) is used as the evaluation metric.
 
 | Model | Input size| MMPose (s/iter) | HRNet (s/iter) | MMPose (mAP) | HRNet (mAP) |
 | :--- | :---------------: | :---------------: |:--------------------: | :----------------------------: | :-----------------: |

diff --git a/docs/getting_started.md b/docs/getting_started.md
@@ -72,18 +72,18 @@ You can use the following commands to test a dataset.
 
 ```shell
 # single-gpu testing
-python tools/test.py ${CONFIG_FILE} ${CHECKPOINT_FILE} [--out ${RESULT_FILE}] [--eval ${EVAL_METRICS}] \
+python tools/test.py ${CONFIG_FILE} ${CHECKPOINT_FILE} [--out ${RESULT_FILE}] [--eval ${EVAL_METRIC}] \
     [--proc_per_gpu ${NUM_PROC_PER_GPU}] [--gpu_collect] [--tmpdir ${TMPDIR}] [--average_clips ${AVG_TYPE}] \
     [--launcher ${JOB_LAUNCHER}] [--local_rank ${LOCAL_RANK}]
 
-python tools/test.py ${CONFIG_FILE} ${CHECKPOINT_FILE} ${GPU_NUM} [--out ${RESULT_FILE}] [--eval ${EVAL_METRICS}] \
+python tools/test.py ${CONFIG_FILE} ${CHECKPOINT_FILE} ${GPU_NUM} [--out ${RESULT_FILE}] [--eval ${EVAL_METRIC}] \
     [--proc_per_gpu ${NUM_PROC_PER_GPU}] [--gpu_collect] [--tmpdir ${TMPDIR}] [--average_clips ${AVG_TYPE}] \
     [--launcher ${JOB_LAUNCHER}] [--local_rank ${LOCAL_RANK}]
 ```
 
 Optional arguments:
 - `RESULT_FILE`: Filename of the output results. If not specified, the results will not be saved to a file.
-- `EVAL_METRICS`: Items to be evaluated on the results. Allowed values depend on the dataset.
+- `EVAL_METRIC`: Items to be evaluated on the results. Allowed values depend on the dataset.
 - `NUM_PROC_PER_GPU`: Number of processes per GPU. If not specified, only one process will be assigned for a single gpu.
 - `--gpu_collect`: If specified, recognition results will be collected using gpu communication. Otherwise, it will save the results on different gpus to `TMPDIR` and collect them by the rank 0 worker.
 - `TMPDIR`: Temporary directory used for collecting results from multiple workers, available when `--gpu_collect` is not specified.

diff --git a/mmpose/core/evaluation/__init__.py b/mmpose/core/evaluation/__init__.py
@@ -1,6 +1,9 @@
-from .acc import keypoints_from_heatmaps, pose_pck_accuracy
+from .bottom_up_eval import (aggregate_results, get_group_preds,
+                             get_multi_stage_outputs)
 from .eval_hooks import DistEvalHook, EvalHook
+from .top_down_eval import keypoints_from_heatmaps, pose_pck_accuracy
 
 __all__ = [
-    'EvalHook', 'DistEvalHook', 'pose_pck_accuracy', 'keypoints_from_heatmaps'
+    'EvalHook', 'DistEvalHook', 'pose_pck_accuracy', 'keypoints_from_heatmaps',
+    'get_group_preds', 'get_multi_stage_outputs', 'aggregate_results'
 ]
diff --git a/mmpose/core/evaluation/bottom_up_eval.py b/mmpose/core/evaluation/bottom_up_eval.py
@@ -0,0 +1,192 @@
+import numpy as np
+import torch
+
+from mmpose.core.post_processing import transform_preds
+
+
+def get_multi_stage_outputs(outputs,
+                            outputs_flip,
+                            num_joints,
+                            with_heatmaps,
+                            with_ae,
+                            tag_per_joint=True,
+                            flip_index=None,
+                            project2image=True,
+                            size_projected=None):
+    """Inference the model to get multi-stage outputs (heatmaps & tags), and
+    resize them to base sizes.
+
+    Args:
+        outputs (list(torch.Tensor)): Outputs of network
+        outputs_flip (list(torch.Tensor)): Flip outputs of network
+        num_joints (int): Number of joints
+        with_heatmaps (list[bool]): Option to output
+            heatmaps for different stages.
+        with_ae (list[bool]): Option to output
+            ae tags for different stages.
+        tag_per_joint (bool): Option to use one tag map per joint.
+        flip_index (list[int]): Keypoint flip index.
+        project2image (bool): Option to resize to base scale.
+        size_projected ([w, h]): Base size of heatmaps.
+
+    Returns:
+        outputs (list(torch.Tensor)): List of simple outputs and
+            flip outputs.
+        heatmaps (torch.Tensor): Multi-stage heatmaps that are resized to
+            the base size.
+        tags (torch.Tensor): Multi-stage tags that are resized to
+            the base size.
+    """
+
+    heatmaps_avg = 0
+    num_heatmaps = 0
+    heatmaps = []
+    tags = []
+
+    flip_test = outputs_flip is not None
+
+    # aggregate heatmaps from different stages
+    for i, output in enumerate(outputs):
+        if i != len(outputs) - 1:
+            output = torch.nn.functional.interpolate(
+                output,
+                size=(outputs[-1].size(2), outputs[-1].size(3)),
+                mode='bilinear',
+                align_corners=False)
+
+        # staring index of the associative embeddings
+        offset_feat = num_joints if with_heatmaps[i] else 0
+
+        if with_heatmaps[i]:
+            heatmaps_avg += output[:, :num_joints]
+            num_heatmaps += 1
+
+        if with_ae[i]:
+            tags.append(output[:, offset_feat:])
+
+    if num_heatmaps > 0:
+        heatmaps.append(heatmaps_avg / num_heatmaps)
+
+    if flip_test and flip_index:
+        # perform flip testing
+        heatmaps_avg = 0
+        num_heatmaps = 0
+
+        for i in range(len(outputs_flip)):
+            output = outputs_flip[i]
+            if i != len(outputs_flip) - 1:
+                output = torch.nn.functional.interpolate(
+                    output,
+                    size=(outputs_flip[-1].size(2), outputs_flip[-1].size(3)),
+                    mode='bilinear',
+                    align_corners=False)
+            output = torch.flip(output, [3])
+            outputs.append(output)
+
+            offset_feat = num_joints if with_heatmaps[i] else 0
+
+            if with_heatmaps[i]:
+                heatmaps_avg += output[:, :num_joints][:, flip_index, :, :]
+                num_heatmaps += 1
+
+            if with_ae[i]:
+                tags.append(output[:, offset_feat:])
+                if tag_per_joint:
+                    tags[-1] = tags[-1][:, flip_index, :, :]
+
+        heatmaps.append(heatmaps_avg / num_heatmaps)
+
+    if project2image and size_projected:
+        heatmaps = [
+            torch.nn.functional.interpolate(
+                hms,
+                size=(size_projected[1], size_projected[0]),
+                mode='bilinear',
+                align_corners=False) for hms in heatmaps
+        ]
+
+        tags = [
+            torch.nn.functional.interpolate(
+                tms,
+                size=(size_projected[1], size_projected[0]),
+                mode='bilinear',
+                align_corners=False) for tms in tags
+        ]
+
+    return outputs, heatmaps, tags
+
+
+def aggregate_results(scale, aggregated_heatmaps, tags_list, heatmaps, tags,
+                      test_scale_factor, project2image, flip_test):
+    """Aggregate multi-scale outputs.
+
+    Note:
+        batch size: N
+        keypoints num : K
+        heatmap width: W
+        heatmap hight: H
+
+    Args:
+        scale (int): current scale
+        aggregated_heatmaps (torch.Tensor | None): Aggregated heatmaps.
+        tags_list (list(torch.Tensor)): Tags list of previous scale.
+        heatmaps (List(torch.Tensor[NxKxWxH])): A batch of heatmaps.
+        tags (List(torch.Tensor[NxKxWxH])): A batch of tag maps.
+        test_scale_factor (List(int)): Multi-scale factor for testing.
+        project2image (bool): Option to resize to base scale.
+        flip_test (bool): Option to use flip test.
+
+    Return:
+        aggregated_heatmaps(torch.Tensor): Heatmaps with multi scale.
+        tags_list(list(torch.Tensor)): Tag list of multi scale.
+    """
+    if scale == 1 or len(test_scale_factor) == 1:
+        if aggregated_heatmaps is not None and not project2image:
+            tags = [
+                torch.nn.functional.interpolate(
+                    tms,
+                    size=(aggregated_heatmaps.size(2),
+                          aggregated_heatmaps.size(3)),
+                    mode='bilinear',
+                    align_corners=False) for tms in tags
+            ]
+        for tms in tags:
+            tags_list.append(torch.unsqueeze(tms, dim=4))
+
+    heatmaps_avg = (heatmaps[0] +
+                    heatmaps[1]) / 2.0 if flip_test else heatmaps[0]
+
+    if aggregated_heatmaps is None:
+        aggregated_heatmaps = heatmaps_avg
+    elif project2image:
+        aggregated_heatmaps += heatmaps_avg
+    else:
+        aggregated_heatmaps += torch.nn.functional.interpolate(
+            heatmaps_avg,
+            size=(aggregated_heatmaps.size(2), aggregated_heatmaps.size(3)),
+            mode='bilinear',
+            align_corners=False)
+
+    return aggregated_heatmaps, tags_list
+
+
+def get_group_preds(grouped_joints, center, scale, heatmap_size):
+    """Transform the grouped joints back to the image.
+
+    Args:
+        grouped_joints (list): Grouped person joints.
+        center (np.ndarray[2, ]): Center of the bounding box (x, y).
+        scale (np.ndarray[2, ]): Scale of the bounding box
+            wrt [width, height].
+        heatmap_size (np.ndarray[2, ]): Size of the destination heatmaps.
+
+    Returns:
+        results (List): List of the pose result for each person.
+    """
+    results = []
+    for person in grouped_joints[0]:
+        joints = np.zeros((person.shape[0], 3))
+        joints = transform_preds(person, center, scale, heatmap_size)
+        results.append(joints)
+
+    return results
diff --git a/mmpose/core/evaluation/acc.py → mmpose/core/evaluation/top_down_eval.py b/mmpose/core/evaluation/acc.py → mmpose/core/evaluation/top_down_eval.py
diff --git a/mmpose/core/post_processing/__init__.py b/mmpose/core/post_processing/__init__.py
@@ -1,10 +1,9 @@
 from .nms import oks_nms, soft_oks_nms
-from .shared_transforms import affine_transform, get_3rd_point, rotate_point
-from .top_down_transforms import (flip_back, fliplr_joints,
-                                  get_affine_transform, transform_preds)
+from .post_transforms import (affine_transform, flip_back, fliplr_joints,
+                              get_affine_transform, rotate_point,
+                              transform_preds)
 
 __all__ = [
-    'oks_nms', 'soft_oks_nms', 'get_3rd_point', 'affine_transform',
-    'rotate_point', 'flip_back', 'fliplr_joints', 'transform_preds',
-    'get_affine_transform'
+    'oks_nms', 'soft_oks_nms', 'affine_transform', 'rotate_point', 'flip_back',
+    'fliplr_joints', 'transform_preds', 'get_affine_transform'
 ]