Harsha909/video-pose-normalization

1# Copyright (c) OpenMMLab. All rights reserved.
2from unittest import TestCase
3
4import numpy as np
5
6from mmpose.visualization import FastVisualizer
7
8
9class TestFastVisualizer(TestCase):
10
11    def setUp(self):
12        self.metainfo = {
13            'keypoint_id2name': {
14                0: 'nose',
15                1: 'left_eye',
16                2: 'right_eye'
17            },
18            'keypoint_name2id': {
19                'nose': 0,
20                'left_eye': 1,
21                'right_eye': 2
22            },
23            'keypoint_colors': np.array([[255, 0, 0], [0, 255, 0], [0, 0,
24                                                                    255]]),
25            'skeleton_links': [(0, 1), (1, 2)],
26            'skeleton_link_colors': np.array([[255, 255, 0], [255, 0, 255]])
27        }
28        self.visualizer = FastVisualizer(self.metainfo)
29
30    def test_init(self):
31        self.assertEqual(self.visualizer.radius, 6)
32        self.assertEqual(self.visualizer.line_width, 3)
33        self.assertEqual(self.visualizer.kpt_thr, 0.3)
34        self.assertEqual(self.visualizer.keypoint_id2name,
35                         self.metainfo['keypoint_id2name'])
36        self.assertEqual(self.visualizer.keypoint_name2id,
37                         self.metainfo['keypoint_name2id'])
38        np.testing.assert_array_equal(self.visualizer.keypoint_colors,
39                                      self.metainfo['keypoint_colors'])
40        self.assertEqual(self.visualizer.skeleton_links,
41                         self.metainfo['skeleton_links'])
42        np.testing.assert_array_equal(self.visualizer.skeleton_link_colors,
43                                      self.metainfo['skeleton_link_colors'])
44
45    def test_draw_pose(self):
46        img = np.zeros((480, 640, 3), dtype=np.uint8)
47        instances = type('Instances', (object, ), {})()
48        instances.keypoints = np.array([[[100, 100], [200, 200], [300, 300]]],
49                                       dtype=np.float32)
50        instances.keypoint_scores = np.array([[0.5, 0.5, 0.5]],
51                                             dtype=np.float32)
52
53        self.visualizer.draw_pose(img, instances)
54
55        # Check if keypoints are drawn
56        self.assertNotEqual(img[100, 100].tolist(), [0, 0, 0])
57        self.assertNotEqual(img[200, 200].tolist(), [0, 0, 0])
58        self.assertNotEqual(img[300, 300].tolist(), [0, 0, 0])
59
60        # Check if skeleton links are drawn
61        self.assertNotEqual(img[150, 150].tolist(), [0, 0, 0])
62        self.assertNotEqual(img[250, 250].tolist(), [0, 0, 0])
63
64    def test_draw_pose_with_none_instances(self):
65        img = np.zeros((480, 640, 3), dtype=np.uint8)
66        instances = None
67
68        self.visualizer.draw_pose(img, instances)
69
70        # Check if the image is still empty (black)
71        self.assertEqual(np.count_nonzero(img), 0)
72