File size: 2,828 Bytes
7bed60d |
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 |
# Copyright (c) OpenMMLab. All rights reserved.
from unittest import TestCase
import numpy as np
from mmpose.visualization import FastVisualizer
class TestFastVisualizer(TestCase):
def setUp(self):
self.metainfo = {
'keypoint_id2name': {
0: 'nose',
1: 'left_eye',
2: 'right_eye'
},
'keypoint_name2id': {
'nose': 0,
'left_eye': 1,
'right_eye': 2
},
'keypoint_colors': np.array([[255, 0, 0], [0, 255, 0], [0, 0,
255]]),
'skeleton_links': [(0, 1), (1, 2)],
'skeleton_link_colors': np.array([[255, 255, 0], [255, 0, 255]])
}
self.visualizer = FastVisualizer(self.metainfo)
def test_init(self):
self.assertEqual(self.visualizer.radius, 6)
self.assertEqual(self.visualizer.line_width, 3)
self.assertEqual(self.visualizer.kpt_thr, 0.3)
self.assertEqual(self.visualizer.keypoint_id2name,
self.metainfo['keypoint_id2name'])
self.assertEqual(self.visualizer.keypoint_name2id,
self.metainfo['keypoint_name2id'])
np.testing.assert_array_equal(self.visualizer.keypoint_colors,
self.metainfo['keypoint_colors'])
self.assertEqual(self.visualizer.skeleton_links,
self.metainfo['skeleton_links'])
np.testing.assert_array_equal(self.visualizer.skeleton_link_colors,
self.metainfo['skeleton_link_colors'])
def test_draw_pose(self):
img = np.zeros((480, 640, 3), dtype=np.uint8)
instances = type('Instances', (object, ), {})()
instances.keypoints = np.array([[[100, 100], [200, 200], [300, 300]]],
dtype=np.float32)
instances.keypoint_scores = np.array([[0.5, 0.5, 0.5]],
dtype=np.float32)
self.visualizer.draw_pose(img, instances)
# Check if keypoints are drawn
self.assertNotEqual(img[100, 100].tolist(), [0, 0, 0])
self.assertNotEqual(img[200, 200].tolist(), [0, 0, 0])
self.assertNotEqual(img[300, 300].tolist(), [0, 0, 0])
# Check if skeleton links are drawn
self.assertNotEqual(img[150, 150].tolist(), [0, 0, 0])
self.assertNotEqual(img[250, 250].tolist(), [0, 0, 0])
def test_draw_pose_with_none_instances(self):
img = np.zeros((480, 640, 3), dtype=np.uint8)
instances = None
self.visualizer.draw_pose(img, instances)
# Check if the image is still empty (black)
self.assertEqual(np.count_nonzero(img), 0)
|