mirror of https://github.com/open-mmlab/mmpose
169 lines
6.1 KiB
Python
169 lines
6.1 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
from unittest import TestCase
|
|
|
|
import numpy as np
|
|
|
|
from mmpose.codecs import DecoupledHeatmap
|
|
from mmpose.registry import KEYPOINT_CODECS
|
|
from mmpose.testing import get_coco_sample
|
|
|
|
|
|
class TestDecoupledHeatmap(TestCase):
|
|
|
|
def setUp(self) -> None:
|
|
pass
|
|
|
|
def _make_multi_instance_data(self, data):
|
|
bbox = data['bbox'].reshape(-1, 2, 2)
|
|
keypoints = data['keypoints']
|
|
keypoints_visible = data['keypoints_visible']
|
|
|
|
keypoints_visible[..., 0] = 0
|
|
|
|
offset = keypoints.max(axis=1, keepdims=True)
|
|
bbox_outside = bbox - offset
|
|
keypoints_outside = keypoints - offset
|
|
keypoints_outside_visible = np.zeros(keypoints_visible.shape)
|
|
|
|
bbox_overlap = bbox.mean(
|
|
axis=1, keepdims=True) + 0.8 * (
|
|
bbox - bbox.mean(axis=1, keepdims=True))
|
|
keypoint_overlap = keypoints.mean(
|
|
axis=1, keepdims=True) + 0.8 * (
|
|
keypoints - keypoints.mean(axis=1, keepdims=True))
|
|
keypoint_overlap_visible = keypoints_visible
|
|
|
|
data['bbox'] = np.concatenate((bbox, bbox_outside, bbox_overlap),
|
|
axis=0)
|
|
data['keypoints'] = np.concatenate(
|
|
(keypoints, keypoints_outside, keypoint_overlap), axis=0)
|
|
data['keypoints_visible'] = np.concatenate(
|
|
(keypoints_visible, keypoints_outside_visible,
|
|
keypoint_overlap_visible),
|
|
axis=0)
|
|
|
|
return data
|
|
|
|
def test_build(self):
|
|
cfg = dict(
|
|
type='DecoupledHeatmap',
|
|
input_size=(512, 512),
|
|
heatmap_size=(128, 128),
|
|
)
|
|
codec = KEYPOINT_CODECS.build(cfg)
|
|
self.assertIsInstance(codec, DecoupledHeatmap)
|
|
|
|
def test_encode(self):
|
|
data = get_coco_sample(img_shape=(512, 512), num_instances=1)
|
|
data['bbox'] = np.tile(data['bbox'], 2).reshape(-1, 4, 2)
|
|
data['bbox'][:, 1:3, 0] = data['bbox'][:, 0:2, 0]
|
|
data = self._make_multi_instance_data(data)
|
|
|
|
codec = DecoupledHeatmap(
|
|
input_size=(512, 512),
|
|
heatmap_size=(128, 128),
|
|
)
|
|
|
|
print(data['bbox'].shape)
|
|
encoded = codec.encode(
|
|
data['keypoints'], data['keypoints_visible'], bbox=data['bbox'])
|
|
|
|
heatmaps = encoded['heatmaps']
|
|
instance_heatmaps = encoded['instance_heatmaps']
|
|
keypoint_weights = encoded['keypoint_weights']
|
|
instance_coords = encoded['instance_coords']
|
|
|
|
self.assertEqual(heatmaps.shape, (18, 128, 128))
|
|
self.assertEqual(keypoint_weights.shape, (2, 17))
|
|
self.assertEqual(instance_heatmaps.shape, (34, 128, 128))
|
|
self.assertEqual(instance_coords.shape, (2, 2))
|
|
|
|
# without bbox
|
|
encoded = codec.encode(
|
|
data['keypoints'], data['keypoints_visible'], bbox=None)
|
|
|
|
heatmaps = encoded['heatmaps']
|
|
instance_heatmaps = encoded['instance_heatmaps']
|
|
keypoint_weights = encoded['keypoint_weights']
|
|
instance_coords = encoded['instance_coords']
|
|
|
|
self.assertEqual(heatmaps.shape, (18, 128, 128))
|
|
self.assertEqual(keypoint_weights.shape, (2, 17))
|
|
self.assertEqual(instance_heatmaps.shape, (34, 128, 128))
|
|
self.assertEqual(instance_coords.shape, (2, 2))
|
|
|
|
# root_type
|
|
with self.assertRaises(ValueError):
|
|
codec = DecoupledHeatmap(
|
|
input_size=(512, 512),
|
|
heatmap_size=(128, 128),
|
|
root_type='box_center',
|
|
)
|
|
encoded = codec.encode(
|
|
data['keypoints'],
|
|
data['keypoints_visible'],
|
|
bbox=data['bbox'])
|
|
|
|
codec = DecoupledHeatmap(
|
|
input_size=(512, 512),
|
|
heatmap_size=(128, 128),
|
|
root_type='bbox_center',
|
|
)
|
|
|
|
encoded = codec.encode(
|
|
data['keypoints'], data['keypoints_visible'], bbox=data['bbox'])
|
|
|
|
heatmaps = encoded['heatmaps']
|
|
instance_heatmaps = encoded['instance_heatmaps']
|
|
keypoint_weights = encoded['keypoint_weights']
|
|
instance_coords = encoded['instance_coords']
|
|
|
|
self.assertEqual(heatmaps.shape, (18, 128, 128))
|
|
self.assertEqual(keypoint_weights.shape, (2, 17))
|
|
self.assertEqual(instance_heatmaps.shape, (34, 128, 128))
|
|
self.assertEqual(instance_coords.shape, (2, 2))
|
|
|
|
def test_decode(self):
|
|
data = get_coco_sample(img_shape=(512, 512), num_instances=2)
|
|
data['bbox'] = np.tile(data['bbox'], 2).reshape(-1, 4, 2)
|
|
data['bbox'][:, 1:3, 0] = data['bbox'][:, 0:2, 0]
|
|
|
|
codec = DecoupledHeatmap(
|
|
input_size=(512, 512),
|
|
heatmap_size=(128, 128),
|
|
)
|
|
|
|
encoded = codec.encode(
|
|
data['keypoints'], data['keypoints_visible'], bbox=data['bbox'])
|
|
instance_heatmaps = encoded['instance_heatmaps'].reshape(
|
|
encoded['instance_coords'].shape[0], -1,
|
|
*encoded['instance_heatmaps'].shape[-2:])
|
|
instance_scores = np.ones(encoded['instance_coords'].shape[0])
|
|
decoded = codec.decode(instance_heatmaps, instance_scores[:, None])
|
|
keypoints, keypoint_scores = decoded
|
|
|
|
self.assertEqual(keypoints.shape, (2, 17, 2))
|
|
self.assertEqual(keypoint_scores.shape, (2, 17))
|
|
|
|
def test_cicular_verification(self):
|
|
data = get_coco_sample(img_shape=(512, 512), num_instances=1)
|
|
data['bbox'] = np.tile(data['bbox'], 2).reshape(-1, 4, 2)
|
|
data['bbox'][:, 1:3, 0] = data['bbox'][:, 0:2, 0]
|
|
|
|
codec = DecoupledHeatmap(
|
|
input_size=(512, 512),
|
|
heatmap_size=(128, 128),
|
|
)
|
|
|
|
encoded = codec.encode(
|
|
data['keypoints'], data['keypoints_visible'], bbox=data['bbox'])
|
|
instance_heatmaps = encoded['instance_heatmaps'].reshape(
|
|
encoded['instance_coords'].shape[0], -1,
|
|
*encoded['instance_heatmaps'].shape[-2:])
|
|
instance_scores = np.ones(encoded['instance_coords'].shape[0])
|
|
decoded = codec.decode(instance_heatmaps, instance_scores[:, None])
|
|
keypoints, _ = decoded
|
|
keypoints += 1.5
|
|
|
|
self.assertTrue(np.allclose(keypoints, data['keypoints'], atol=5.))
|