def test_unqualified_detection(self):
detection = text_format.Parse('location_data {format: GLOBAL}',
detection_pb2.Detection())
image = np.arange(27, dtype=np.uint8).reshape(3, 3, 3)
with self.assertRaisesRegex(ValueError,
'LocationData must be relative'):
drawing_utils.draw_detection(image, detection)
def _annotate(self, frame: np.ndarray, results: NamedTuple, idx: int):
for detection in results.detections:
mp_drawing.draw_detection(frame, detection)
path = os.path.join(
tempfile.gettempdir(),
self.id().split('.')[-1] + '_frame_{}.png'.format(idx))
cv2.imwrite(path, frame)
def test_invalid_input_image(self):
image = np.arange(18, dtype=np.uint8).reshape(3, 3, 2)
with self.assertRaisesRegex(
ValueError,
'Input image must contain three channel rgb data.'):
drawing_utils.draw_landmarks(image,
landmark_pb2.NormalizedLandmarkList())
with self.assertRaisesRegex(
ValueError,
'Input image must contain three channel rgb data.'):
drawing_utils.draw_detection(image, detection_pb2.Detection())
def test_draw_bboxs_only(self):
detection = text_format.Parse(
'location_data {'
' format: RELATIVE_BOUNDING_BOX'
' relative_bounding_box {xmin: 0 ymin: 0 width: 1 height: 1}}',
detection_pb2.Detection())
image = np.zeros((100, 100, 3), np.uint8)
expected_result = np.copy(image)
cv2.rectangle(expected_result, (0, 0), (99, 99),
DEFAULT_BBOX_DRAWING_SPEC.color,
DEFAULT_BBOX_DRAWING_SPEC.thickness)
drawing_utils.draw_detection(image, detection)
np.testing.assert_array_equal(image, expected_result)
def test_draw_keypoints_only(self):
detection = text_format.Parse(
'location_data {'
' format: RELATIVE_BOUNDING_BOX'
' relative_keypoints {x: 0 y: 1}'
' relative_keypoints {x: 1 y: 0}}', detection_pb2.Detection())
image = np.zeros((100, 100, 3), np.uint8)
expected_result = np.copy(image)
cv2.circle(expected_result, (0, 99),
DEFAULT_CIRCLE_DRAWING_SPEC.circle_radius,
DEFAULT_CIRCLE_DRAWING_SPEC.color,
DEFAULT_CIRCLE_DRAWING_SPEC.thickness)
cv2.circle(expected_result, (99, 0),
DEFAULT_CIRCLE_DRAWING_SPEC.circle_radius,
DEFAULT_CIRCLE_DRAWING_SPEC.color,
DEFAULT_CIRCLE_DRAWING_SPEC.thickness)
drawing_utils.draw_detection(image, detection)
np.testing.assert_array_equal(image, expected_result)