Harsha909/video-pose-normalization

1# Copyright (c) OpenMMLab. All rights reserved.
2from itertools import product
3from unittest import TestCase
4
5import numpy as np
6import torch
7from munkres import Munkres
8
9from mmpose.codecs import AssociativeEmbedding
10from mmpose.registry import KEYPOINT_CODECS
11from mmpose.testing import get_coco_sample
12
13
14class TestAssociativeEmbedding(TestCase):
15
16    def setUp(self) -> None:
17        self.decode_keypoint_order = [
18            0, 1, 2, 3, 4, 5, 6, 11, 12, 7, 8, 9, 10, 13, 14, 15, 16
19        ]
20
21    def test_build(self):
22        cfg = dict(
23            type='AssociativeEmbedding',
24            input_size=(256, 256),
25            heatmap_size=(64, 64),
26            use_udp=False,
27            decode_keypoint_order=self.decode_keypoint_order,
28        )
29        codec = KEYPOINT_CODECS.build(cfg)
30        self.assertIsInstance(codec, AssociativeEmbedding)
31
32    def test_encode(self):
33        data = get_coco_sample(img_shape=(256, 256), num_instances=1)
34
35        # w/o UDP
36        codec = AssociativeEmbedding(
37            input_size=(256, 256),
38            heatmap_size=(64, 64),
39            use_udp=False,
40            decode_keypoint_order=self.decode_keypoint_order)
41
42        encoded = codec.encode(data['keypoints'], data['keypoints_visible'])
43
44        heatmaps = encoded['heatmaps']
45        keypoint_indices = encoded['keypoint_indices']
46        keypoint_weights = encoded['keypoint_weights']
47
48        self.assertEqual(heatmaps.shape, (17, 64, 64))
49        self.assertEqual(keypoint_indices.shape, (1, 17, 2))
50        self.assertEqual(keypoint_weights.shape, (1, 17))
51
52        for k in range(heatmaps.shape[0]):
53            index_expected = np.argmax(heatmaps[k])
54            index_encoded = keypoint_indices[0, k, 0]
55            self.assertEqual(index_expected, index_encoded)
56
57        # w/ UDP
58        codec = AssociativeEmbedding(
59            input_size=(256, 256),
60            heatmap_size=(64, 64),
61            use_udp=True,
62            decode_keypoint_order=self.decode_keypoint_order)
63
64        encoded = codec.encode(data['keypoints'], data['keypoints_visible'])
65
66        heatmaps = encoded['heatmaps']
67        keypoint_indices = encoded['keypoint_indices']
68        keypoint_weights = encoded['keypoint_weights']
69
70        self.assertEqual(heatmaps.shape, (17, 64, 64))
71        self.assertEqual(keypoint_indices.shape, (1, 17, 2))
72        self.assertEqual(keypoint_weights.shape, (1, 17))
73
74        for k in range(heatmaps.shape[0]):
75            index_expected = np.argmax(heatmaps[k])
76            index_encoded = keypoint_indices[0, k, 0]
77            self.assertEqual(index_expected, index_encoded)
78
79    def _get_tags(self,
80                  heatmaps,
81                  keypoint_indices,
82                  tag_per_keypoint: bool,
83                  tag_dim: int = 1):
84
85        K, H, W = heatmaps.shape
86        N = keypoint_indices.shape[0]
87
88        if tag_per_keypoint:
89            tags = np.zeros((K * tag_dim, H, W), dtype=np.float32)
90        else:
91            tags = np.zeros((tag_dim, H, W), dtype=np.float32)
92
93        for n, k in product(range(N), range(K)):
94            y, x = np.unravel_index(keypoint_indices[n, k, 0], (H, W))
95            if tag_per_keypoint:
96                tags[k::K, y, x] = n
97            else:
98                tags[:, y, x] = n
99
100        return tags
101
102    def _sort_preds(self, keypoints_pred, scores_pred, keypoints_gt):
103        """Sort multi-instance predictions to best match the ground-truth.
104
105        Args:
106            keypoints_pred (np.ndarray): predictions in shape (N, K, D)
107            scores (np.ndarray): predictions in shape (N, K)
108            keypoints_gt (np.ndarray): ground-truth in shape (N, K, D)
109
110        Returns:
111            np.ndarray: Sorted predictions
112        """
113        assert keypoints_gt.shape == keypoints_pred.shape
114        costs = np.linalg.norm(
115            keypoints_gt[None] - keypoints_pred[:, None], ord=2,
116            axis=3).mean(axis=2)
117        match = Munkres().compute(costs)
118        keypoints_pred_sorted = np.zeros_like(keypoints_pred)
119        scores_pred_sorted = np.zeros_like(scores_pred)
120        for i, j in match:
121            keypoints_pred_sorted[i] = keypoints_pred[j]
122            scores_pred_sorted[i] = scores_pred[j]
123
124        return keypoints_pred_sorted, scores_pred_sorted
125
126    def test_decode(self):
127        data = get_coco_sample(
128            img_shape=(256, 256), num_instances=2, non_occlusion=True)
129
130        # w/o UDP
131        codec = AssociativeEmbedding(
132            input_size=(256, 256),
133            heatmap_size=(64, 64),
134            use_udp=False,
135            decode_keypoint_order=self.decode_keypoint_order)
136
137        encoded = codec.encode(data['keypoints'], data['keypoints_visible'])
138
139        heatmaps = encoded['heatmaps']
140        keypoint_indices = encoded['keypoint_indices']
141
142        tags = self._get_tags(
143            heatmaps, keypoint_indices, tag_per_keypoint=True)
144
145        # to Tensor
146        batch_heatmaps = torch.from_numpy(heatmaps[None])
147        batch_tags = torch.from_numpy(tags[None])
148
149        batch_keypoints, batch_keypoint_scores = codec.batch_decode(
150            batch_heatmaps, batch_tags)
151
152        self.assertIsInstance(batch_keypoints, list)
153        self.assertIsInstance(batch_keypoint_scores, list)
154        self.assertEqual(len(batch_keypoints), 1)
155        self.assertEqual(len(batch_keypoint_scores), 1)
156
157        keypoints, scores = self._sort_preds(batch_keypoints[0],
158                                             batch_keypoint_scores[0],
159                                             data['keypoints'])
160
161        self.assertIsInstance(keypoints, np.ndarray)
162        self.assertIsInstance(scores, np.ndarray)
163        self.assertEqual(keypoints.shape, (2, 17, 2))
164        self.assertEqual(scores.shape, (2, 17))
165
166        self.assertTrue(np.allclose(keypoints, data['keypoints'], atol=4.0))
167
168        # w/o UDP, tag_imd=2
169        codec = AssociativeEmbedding(
170            input_size=(256, 256),
171            heatmap_size=(64, 64),
172            use_udp=False,
173            decode_keypoint_order=self.decode_keypoint_order)
174
175        encoded = codec.encode(data['keypoints'], data['keypoints_visible'])
176
177        heatmaps = encoded['heatmaps']
178        keypoint_indices = encoded['keypoint_indices']
179
180        tags = self._get_tags(
181            heatmaps, keypoint_indices, tag_per_keypoint=True, tag_dim=2)
182
183        # to Tensor
184        batch_heatmaps = torch.from_numpy(heatmaps[None])
185        batch_tags = torch.from_numpy(tags[None])
186
187        batch_keypoints, batch_keypoint_scores = codec.batch_decode(
188            batch_heatmaps, batch_tags)
189
190        self.assertIsInstance(batch_keypoints, list)
191        self.assertIsInstance(batch_keypoint_scores, list)
192        self.assertEqual(len(batch_keypoints), 1)
193        self.assertEqual(len(batch_keypoint_scores), 1)
194
195        keypoints, scores = self._sort_preds(batch_keypoints[0],
196                                             batch_keypoint_scores[0],
197                                             data['keypoints'])
198
199        self.assertIsInstance(keypoints, np.ndarray)
200        self.assertIsInstance(scores, np.ndarray)
201        self.assertEqual(keypoints.shape, (2, 17, 2))
202        self.assertEqual(scores.shape, (2, 17))
203
204        self.assertTrue(np.allclose(keypoints, data['keypoints'], atol=4.0))
205
206        # w/ UDP
207        codec = AssociativeEmbedding(
208            input_size=(256, 256),
209            heatmap_size=(64, 64),
210            use_udp=True,
211            decode_keypoint_order=self.decode_keypoint_order)
212
213        encoded = codec.encode(data['keypoints'], data['keypoints_visible'])
214
215        heatmaps = encoded['heatmaps']
216        keypoint_indices = encoded['keypoint_indices']
217
218        tags = self._get_tags(
219            heatmaps, keypoint_indices, tag_per_keypoint=True)
220
221        # to Tensor
222        batch_heatmaps = torch.from_numpy(heatmaps[None])
223        batch_tags = torch.from_numpy(tags[None])
224
225        batch_keypoints, batch_keypoint_scores = codec.batch_decode(
226            batch_heatmaps, batch_tags)
227
228        self.assertIsInstance(batch_keypoints, list)
229        self.assertIsInstance(batch_keypoint_scores, list)
230        self.assertEqual(len(batch_keypoints), 1)
231        self.assertEqual(len(batch_keypoint_scores), 1)
232
233        keypoints, scores = self._sort_preds(batch_keypoints[0],
234                                             batch_keypoint_scores[0],
235                                             data['keypoints'])
236
237        self.assertIsInstance(keypoints, np.ndarray)
238        self.assertIsInstance(scores, np.ndarray)
239        self.assertEqual(keypoints.shape, (2, 17, 2))
240        self.assertEqual(scores.shape, (2, 17))
241
242        self.assertTrue(np.allclose(keypoints, data['keypoints'], atol=4.0))
243