# Copyright 2020 The MediaPipe Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for mediapipe.python.solutions.drawing_utils."""

from absl.testing import absltest
from absl.testing import parameterized
import cv2
import numpy as np

from google.protobuf import text_format

from mediapipe.framework.formats import detection_pb2
from mediapipe.framework.formats import landmark_pb2
from mediapipe.python.solutions import drawing_utils

DEFAULT_BBOX_DRAWING_SPEC = drawing_utils.DrawingSpec()
DEFAULT_CONNECTION_DRAWING_SPEC = drawing_utils.DrawingSpec()
DEFAULT_CIRCLE_DRAWING_SPEC = drawing_utils.DrawingSpec(
    color=drawing_utils.RED_COLOR)
DEFAULT_AXIS_DRAWING_SPEC = drawing_utils.DrawingSpec()
DEFAULT_CYCLE_BORDER_COLOR = (224, 224, 224)


class DrawingUtilTest(parameterized.TestCase):

  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 bgr data.'):
      drawing_utils.draw_landmarks(image, landmark_pb2.NormalizedLandmarkList())
    with self.assertRaisesRegex(
        ValueError, 'Input image must contain three channel bgr data.'):
      drawing_utils.draw_detection(image, detection_pb2.Detection())
    with self.assertRaisesRegex(
        ValueError, 'Input image must contain three channel bgr data.'):
      rotation = np.eye(3, dtype=np.float32)
      translation = np.array([0., 0., 1.])
      drawing_utils.draw_axis(image, rotation, translation)

  def test_invalid_connection(self):
    landmark_list = text_format.Parse(
        'landmark {x: 0.5 y: 0.5} landmark {x: 0.2 y: 0.2}',
        landmark_pb2.NormalizedLandmarkList())
    image = np.arange(27, dtype=np.uint8).reshape(3, 3, 3)
    with self.assertRaisesRegex(ValueError, 'Landmark index is out of range.'):
      drawing_utils.draw_landmarks(image, landmark_list, [(0, 2)])

  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 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)

  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)

  @parameterized.named_parameters(
      ('landmark_list_has_only_one_element', 'landmark {x: 0.1 y: 0.1}'),
      ('second_landmark_is_invisible',
       'landmark {x: 0.1 y: 0.1} landmark {x: 0.5 y: 0.5 visibility: 0.0}'))
  def test_draw_single_landmark_point(self, landmark_list_text):
    landmark_list = text_format.Parse(landmark_list_text,
                                      landmark_pb2.NormalizedLandmarkList())
    image = np.zeros((100, 100, 3), np.uint8)
    expected_result = np.copy(image)
    cv2.circle(expected_result, (10, 10),
               DEFAULT_CIRCLE_DRAWING_SPEC.circle_radius + 1,
               DEFAULT_CYCLE_BORDER_COLOR,
               DEFAULT_CIRCLE_DRAWING_SPEC.thickness)
    cv2.circle(expected_result, (10, 10),
               DEFAULT_CIRCLE_DRAWING_SPEC.circle_radius,
               DEFAULT_CIRCLE_DRAWING_SPEC.color,
               DEFAULT_CIRCLE_DRAWING_SPEC.thickness)
    drawing_utils.draw_landmarks(image, landmark_list)
    np.testing.assert_array_equal(image, expected_result)

  @parameterized.named_parameters(
      ('landmarks_have_x_and_y_only',
       'landmark {x: 0.1 y: 0.5} landmark {x: 0.5 y: 0.1}'),
      ('landmark_zero_visibility_and_presence',
       'landmark {x: 0.1 y: 0.5 presence: 0.5}'
       'landmark {x: 0.5 y: 0.1 visibility: 0.5}'))
  def test_draw_landmarks_and_connections(self, landmark_list_text):
    landmark_list = text_format.Parse(landmark_list_text,
                                      landmark_pb2.NormalizedLandmarkList())
    image = np.zeros((100, 100, 3), np.uint8)
    expected_result = np.copy(image)
    start_point = (10, 50)
    end_point = (50, 10)
    cv2.line(expected_result, start_point, end_point,
             DEFAULT_CONNECTION_DRAWING_SPEC.color,
             DEFAULT_CONNECTION_DRAWING_SPEC.thickness)
    cv2.circle(expected_result, start_point,
               DEFAULT_CIRCLE_DRAWING_SPEC.circle_radius + 1,
               DEFAULT_CYCLE_BORDER_COLOR,
               DEFAULT_CIRCLE_DRAWING_SPEC.thickness)
    cv2.circle(expected_result, end_point,
               DEFAULT_CIRCLE_DRAWING_SPEC.circle_radius + 1,
               DEFAULT_CYCLE_BORDER_COLOR,
               DEFAULT_CIRCLE_DRAWING_SPEC.thickness)
    cv2.circle(expected_result, start_point,
               DEFAULT_CIRCLE_DRAWING_SPEC.circle_radius,
               DEFAULT_CIRCLE_DRAWING_SPEC.color,
               DEFAULT_CIRCLE_DRAWING_SPEC.thickness)
    cv2.circle(expected_result, end_point,
               DEFAULT_CIRCLE_DRAWING_SPEC.circle_radius,
               DEFAULT_CIRCLE_DRAWING_SPEC.color,
               DEFAULT_CIRCLE_DRAWING_SPEC.thickness)
    drawing_utils.draw_landmarks(
        image=image, landmark_list=landmark_list, connections=[(0, 1)])
    np.testing.assert_array_equal(image, expected_result)

  def test_draw_axis(self):
    image = np.zeros((100, 100, 3), np.uint8)
    expected_result = np.copy(image)
    origin = (50, 50)
    x_axis = (75, 50)
    y_axis = (50, 22)
    z_axis = (50, 77)
    cv2.arrowedLine(expected_result, origin, x_axis, drawing_utils.RED_COLOR,
                    DEFAULT_AXIS_DRAWING_SPEC.thickness)
    cv2.arrowedLine(expected_result, origin, y_axis, drawing_utils.GREEN_COLOR,
                    DEFAULT_AXIS_DRAWING_SPEC.thickness)
    cv2.arrowedLine(expected_result, origin, z_axis, drawing_utils.BLUE_COLOR,
                    DEFAULT_AXIS_DRAWING_SPEC.thickness)
    r = np.sqrt(2.) / 2.
    rotation = np.array([[1., 0., 0.], [0., r, -r], [0., r, r]])
    translation = np.array([0, 0, -0.2])
    drawing_utils.draw_axis(image, rotation, translation)
    np.testing.assert_array_equal(image, expected_result)

  def test_draw_axis_zero_translation(self):
    image = np.zeros((100, 100, 3), np.uint8)
    expected_result = np.copy(image)
    origin = (50, 50)
    x_axis = (0, 50)
    y_axis = (50, 100)
    z_axis = (50, 50)
    cv2.arrowedLine(expected_result, origin, x_axis, drawing_utils.RED_COLOR,
                    DEFAULT_AXIS_DRAWING_SPEC.thickness)
    cv2.arrowedLine(expected_result, origin, y_axis, drawing_utils.GREEN_COLOR,
                    DEFAULT_AXIS_DRAWING_SPEC.thickness)
    cv2.arrowedLine(expected_result, origin, z_axis, drawing_utils.BLUE_COLOR,
                    DEFAULT_AXIS_DRAWING_SPEC.thickness)
    rotation = np.eye(3, dtype=np.float32)
    translation = np.zeros((3,), dtype=np.float32)
    drawing_utils.draw_axis(image, rotation, translation)
    np.testing.assert_array_equal(image, expected_result)

  def test_min_and_max_coordinate_values(self):
    landmark_list = text_format.Parse(
        'landmark {x: 0.0 y: 1.0}'
        'landmark {x: 1.0 y: 0.0}', landmark_pb2.NormalizedLandmarkList())
    image = np.zeros((100, 100, 3), np.uint8)
    expected_result = np.copy(image)
    start_point = (0, 99)
    end_point = (99, 0)
    cv2.line(expected_result, start_point, end_point,
             DEFAULT_CONNECTION_DRAWING_SPEC.color,
             DEFAULT_CONNECTION_DRAWING_SPEC.thickness)
    cv2.circle(expected_result, start_point,
               DEFAULT_CIRCLE_DRAWING_SPEC.circle_radius + 1,
               DEFAULT_CYCLE_BORDER_COLOR,
               DEFAULT_CIRCLE_DRAWING_SPEC.thickness)
    cv2.circle(expected_result, end_point,
               DEFAULT_CIRCLE_DRAWING_SPEC.circle_radius + 1,
               DEFAULT_CYCLE_BORDER_COLOR,
               DEFAULT_CIRCLE_DRAWING_SPEC.thickness)
    cv2.circle(expected_result, start_point,
               DEFAULT_CIRCLE_DRAWING_SPEC.circle_radius,
               DEFAULT_CIRCLE_DRAWING_SPEC.color,
               DEFAULT_CIRCLE_DRAWING_SPEC.thickness)
    cv2.circle(expected_result, end_point,
               DEFAULT_CIRCLE_DRAWING_SPEC.circle_radius,
               DEFAULT_CIRCLE_DRAWING_SPEC.color,
               DEFAULT_CIRCLE_DRAWING_SPEC.thickness)
    drawing_utils.draw_landmarks(
        image=image, landmark_list=landmark_list, connections=[(0, 1)])
    np.testing.assert_array_equal(image, expected_result)

  def test_drawing_spec(self):
    landmark_list = text_format.Parse(
        'landmark {x: 0.1 y: 0.1}'
        'landmark {x: 0.8 y: 0.8}', landmark_pb2.NormalizedLandmarkList())
    image = np.zeros((100, 100, 3), np.uint8)
    landmark_drawing_spec = drawing_utils.DrawingSpec(
        color=(0, 0, 255), thickness=5)
    connection_drawing_spec = drawing_utils.DrawingSpec(
        color=(255, 0, 0), thickness=3)
    expected_result = np.copy(image)
    start_point = (10, 10)
    end_point = (80, 80)
    cv2.line(expected_result, start_point, end_point,
             connection_drawing_spec.color, connection_drawing_spec.thickness)
    cv2.circle(expected_result, start_point,
               landmark_drawing_spec.circle_radius + 1,
               DEFAULT_CYCLE_BORDER_COLOR, landmark_drawing_spec.thickness)
    cv2.circle(expected_result, end_point,
               landmark_drawing_spec.circle_radius + 1,
               DEFAULT_CYCLE_BORDER_COLOR, landmark_drawing_spec.thickness)
    cv2.circle(expected_result, start_point,
               landmark_drawing_spec.circle_radius, landmark_drawing_spec.color,
               landmark_drawing_spec.thickness)
    cv2.circle(expected_result, end_point, landmark_drawing_spec.circle_radius,
               landmark_drawing_spec.color, landmark_drawing_spec.thickness)
    drawing_utils.draw_landmarks(
        image=image,
        landmark_list=landmark_list,
        connections=[(0, 1)],
        landmark_drawing_spec=landmark_drawing_spec,
        connection_drawing_spec=connection_drawing_spec)
    np.testing.assert_array_equal(image, expected_result)


if __name__ == '__main__':
  absltest.main()
