File size: 3,412 Bytes
152f0f2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 |
# Copyright (c) OpenMMLab. All rights reserved.
from unittest import TestCase
from mmpose.datasets.transforms import KeypointConverter
from mmpose.testing import get_coco_sample
class TestKeypointConverter(TestCase):
def setUp(self):
# prepare dummy bottom-up data sample with COCO metainfo
self.data_info = get_coco_sample(
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())
# 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:
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())
# 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())
|