Project import generated by Copybara.

GitOrigin-RevId: f7d09ed033907b893638a8eb4148efa11c0f09a6

mediapipe/python/solutions/drawing_utils.py

@@ -0,0 +1,114 @@
+# 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.
+
+# Lint as: python3
+"""MediaPipe solution drawing utils."""
+
+import math
+from typing import List, Tuple, Union
+
+import cv2
+import dataclasses
+import numpy as np
+
+from mediapipe.framework.formats import landmark_pb2
+
+RGB_CHANNELS = 3
+RED_COLOR = (0, 0, 255)
+
+
+@dataclasses.dataclass
+class DrawingSpec:
+  # Color for drawing the annotation. Default to the green color.
+  color: Tuple[int, int, int] = (0, 255, 0)
+  # Thickness for drawing the annotation. Default to 2 pixels.
+  thickness: int = 2
+  # Circle radius. Default to 2 pixels.
+  circle_radius: int = 2
+
+
+def _normalized_to_pixel_coordinates(
+    normalized_x: float, normalized_y: float, image_width: int,
+    image_height: int) -> Union[None, Tuple[int, int]]:
+  """Converts normalized value pair to pixel coordinates."""
+
+  # Checks if the float value is between 0 and 1.
+  def is_valid_normalized_value(value: float) -> bool:
+    return (value > 0 or math.isclose(0, value)) and (value < 1 or
+                                                      math.isclose(1, value))
+
+  if not (is_valid_normalized_value(normalized_x) and
+          is_valid_normalized_value(normalized_y)):
+    # TODO: Draw coordinates even if it's outside of the image bounds.
+    return None
+  x_px = min(math.floor(normalized_x * image_width), image_width - 1)
+  y_px = min(math.floor(normalized_y * image_height), image_height - 1)
+  return x_px, y_px
+
+
+def draw_landmarks(
+    image: np.ndarray,
+    landmark_list: landmark_pb2.NormalizedLandmarkList,
+    connections: List[Tuple[int, int]] = None,
+    landmark_drawing_spec: DrawingSpec = DrawingSpec(color=RED_COLOR),
+    connection_drawing_spec: DrawingSpec = DrawingSpec()):
+  """Draws the landmarks and the connections on the image.
+
+  Args:
+    image: A three channel RGB image represented as numpy ndarray.
+    landmark_list: A normalized landmark list proto message to be annotated on
+      the image.
+    connections: A list of landmark index tuples that specifies how landmarks to
+      be connected in the drawing.
+    landmark_drawing_spec: A DrawingSpec object that specifies the landmarks'
+      drawing settings such as color, line thickness, and circle radius.
+    connection_drawing_spec: A DrawingSpec object that specifies the
+      connections' drawing settings such as color and line thickness.
+
+  Raises:
+    ValueError: If one of the followings:
+      a) If the input image is not three channel RGB.
+      b) If any connetions contain invalid landmark index.
+  """
+  if not landmark_list:
+    return
+  if image.shape[2] != RGB_CHANNELS:
+    raise ValueError('Input image must contain three channel rgb data.')
+  image_rows, image_cols, _ = image.shape
+  idx_to_coordinates = {}
+  for idx, landmark in enumerate(landmark_list.landmark):
+    if landmark.visibility < 0 or landmark.presence < 0:
+      continue
+    landmark_px = _normalized_to_pixel_coordinates(landmark.x, landmark.y,
+                                                   image_cols, image_rows)
+    if landmark_px:
+      idx_to_coordinates[idx] = landmark_px
+  if connections:
+    num_landmarks = len(landmark_list.landmark)
+    # Draws the connections if the start and end landmarks are both visible.
+    for connection in connections:
+      start_idx = connection[0]
+      end_idx = connection[1]
+      if not (0 <= start_idx < num_landmarks and 0 <= end_idx < num_landmarks):
+        raise ValueError(f'Landmark index is out of range. Invalid connection '
+                         f'from landmark #{start_idx} to landmark #{end_idx}.')
+      if start_idx in idx_to_coordinates and end_idx in idx_to_coordinates:
+        cv2.line(image, idx_to_coordinates[start_idx],
+                 idx_to_coordinates[end_idx], connection_drawing_spec.color,
+                 connection_drawing_spec.thickness)
+  # Draws landmark points after finishing the connection lines, which is
+  # aesthetically better.
+  for landmark_px in idx_to_coordinates.values():
+    cv2.circle(image, landmark_px, landmark_drawing_spec.circle_radius,
+               landmark_drawing_spec.color, landmark_drawing_spec.thickness)