| import unittest |
| import numpy as np |
|
|
| import importlib |
| utils = importlib.import_module('extensions.sd-webui-controlnet.tests.utils', 'utils') |
| utils.setup_test_env() |
|
|
| from annotator.openpose.util import faceDetect, handDetect |
| from annotator.openpose.body import Keypoint, BodyResult |
|
|
| class TestFaceDetect(unittest.TestCase): |
| def test_no_faces(self): |
| oriImg = np.zeros((100, 100, 3), dtype=np.uint8) |
| body = BodyResult([None] * 18, total_score=3, total_parts=0) |
| expected_result = None |
| result = faceDetect(body, oriImg) |
|
|
| self.assertEqual(result, expected_result) |
|
|
| def test_single_face(self): |
| body = BodyResult([ |
| Keypoint(50, 50), |
| *([None] * 13), |
| Keypoint(30, 40), |
| Keypoint(70, 40), |
| Keypoint(20, 50), |
| Keypoint(80, 50), |
| ], total_score=2, total_parts=5) |
|
|
| oriImg = np.zeros((100, 100, 3), dtype=np.uint8) |
|
|
| expected_result = (0, 0, 120) |
| result = faceDetect(body, oriImg) |
|
|
| self.assertEqual(result, expected_result) |
|
|
| class TestHandDetect(unittest.TestCase): |
| def test_no_hands(self): |
| oriImg = np.zeros((100, 100, 3), dtype=np.uint8) |
| body = BodyResult([None] * 18, total_score=3, total_parts=0) |
| expected_result = [] |
| result = handDetect(body, oriImg) |
|
|
| self.assertEqual(result, expected_result) |
|
|
| def test_single_left_hand(self): |
| oriImg = np.zeros((100, 100, 3), dtype=np.uint8) |
|
|
| body = BodyResult([ |
| None, None, None, None, None, |
| Keypoint(20, 20), |
| Keypoint(40, 30), |
| Keypoint(60, 40), |
| *([None] * 8), |
| Keypoint(20, 60), |
| Keypoint(40, 70), |
| Keypoint(60, 80) |
| ], total_score=3, total_parts=0.5) |
|
|
| expected_result = [(49, 26, 33, True)] |
| result = handDetect(body, oriImg) |
|
|
| self.assertEqual(result, expected_result) |
|
|
| def test_single_right_hand(self): |
| oriImg = np.zeros((100, 100, 3), dtype=np.uint8) |
|
|
| body = BodyResult([ |
| None, None, |
| Keypoint(20, 20), |
| Keypoint(40, 30), |
| Keypoint(60, 40), |
| *([None] * 11), |
| Keypoint(20, 60), |
| Keypoint(40, 70), |
| Keypoint(60, 80) |
| ], total_score=3, total_parts=0.5) |
|
|
| expected_result = [(49, 26, 33, False)] |
| result = handDetect(body, oriImg) |
|
|
| self.assertEqual(result, expected_result) |
|
|
| def test_multiple_hands(self): |
| body = BodyResult([ |
| Keypoint(20, 20), |
| Keypoint(40, 30), |
| Keypoint(60, 40), |
| Keypoint(20, 60), |
| Keypoint(40, 70), |
| Keypoint(60, 80), |
| Keypoint(10, 10), |
| Keypoint(30, 20), |
| Keypoint(50, 30), |
| Keypoint(10, 50), |
| Keypoint(30, 60), |
| Keypoint(50, 70), |
| *([None] * 6), |
| ], total_score=3, total_parts=0.5) |
|
|
| oriImg = np.zeros((100, 100, 3), dtype=np.uint8) |
|
|
| expected_result = [(0, 0, 100, True), (16, 43, 56, False)] |
| result = handDetect(body, oriImg) |
| self.assertEqual(result, expected_result) |
|
|
|
|
| if __name__ == '__main__': |
| unittest.main() |
|
|
