File size: 2,500 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 |
# Copyright (c) OpenMMLab. All rights reserved.
from unittest import TestCase
import numpy as np
import torch
from mmengine.structures import PixelData
from mmpose.structures import MultilevelPixelData
class TestMultilevelPixelData(TestCase):
def get_multi_level_pixel_data(self):
metainfo = dict(num_keypoints=17)
sizes = [(64, 48), (32, 24), (16, 12)]
heatmaps = [np.random.rand(17, h, w) for h, w in sizes]
masks = [torch.rand(1, h, w) for h, w in sizes]
data = MultilevelPixelData(
metainfo=metainfo, heatmaps=heatmaps, masks=masks)
return data
def test_init(self):
data = self.get_multi_level_pixel_data()
self.assertIn('num_keypoints', data)
self.assertTrue(data.nlevel == 3)
self.assertTrue(data.shape == ((64, 48), (32, 24), (16, 12)))
self.assertTrue(isinstance(data[0], PixelData))
def test_setter(self):
# test `set_field`
data = self.get_multi_level_pixel_data()
sizes = [(64, 48), (32, 24), (16, 8)]
offset_maps = [torch.rand(2, h, w) for h, w in sizes]
data.offset_maps = offset_maps
# test `to_tensor`
data = self.get_multi_level_pixel_data()
self.assertTrue(isinstance(data[0].heatmaps, np.ndarray))
data = data.to_tensor()
self.assertTrue(isinstance(data[0].heatmaps, torch.Tensor))
# test `cpu`
data = self.get_multi_level_pixel_data()
self.assertTrue(isinstance(data[0].heatmaps, np.ndarray))
self.assertTrue(isinstance(data[0].masks, torch.Tensor))
self.assertTrue(data[0].masks.device.type == 'cpu')
data = data.cpu()
self.assertTrue(isinstance(data[0].heatmaps, np.ndarray))
self.assertTrue(data[0].masks.device.type == 'cpu')
# test `to`
data = self.get_multi_level_pixel_data()
self.assertTrue(data[0].masks.device.type == 'cpu')
data = data.to('cpu')
self.assertTrue(data[0].masks.device.type == 'cpu')
# test `numpy`
data = self.get_multi_level_pixel_data()
self.assertTrue(isinstance(data[0].masks, torch.Tensor))
data = data.numpy()
self.assertTrue(isinstance(data[0].masks, np.ndarray))
def test_deleter(self):
data = self.get_multi_level_pixel_data()
for key in ['heatmaps', 'masks']:
self.assertIn(key, data)
exec(f'del data.{key}')
self.assertNotIn(key, data)
|