Spaces:
Sleeping
Sleeping
File size: 5,939 Bytes
31cf66c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 |
# @title Default title text
from mediapipe.framework.formats import detection_pb2
from mediapipe.framework.formats import location_data_pb2
from mediapipe.framework.formats import landmark_pb2
from mediapipe.python.solutions.drawing_utils import DrawingSpec, _normalized_to_pixel_coordinates
import math
from typing import List, Mapping, Optional, Tuple, Union
import numpy as np
PRESENCE_THRESHOLD = 0.5
RGB_CHANNELS = 3
BLACK_COLOR = (0, 0, 0)
RED_COLOR = (0, 0, 255)
GREEN_COLOR = (0, 128, 0)
BLUE_COLOR = (255, 0, 0)
VISIBILITY_THRESHOLD = 0.5
def draw_landmarks(
image: np.ndarray,
landmark_list: landmark_pb2.NormalizedLandmarkList,
connections: Optional[List[Tuple[int, int]]] = None,
landmark_drawing_spec: Union[DrawingSpec,
Mapping[int, DrawingSpec]] = DrawingSpec(
color=RED_COLOR),
connection_drawing_spec: Union[DrawingSpec,
Mapping[Tuple[int, int],
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.HasField('visibility') and
landmark.visibility < VISIBILITY_THRESHOLD) or
(landmark.HasField('presence') and
landmark.presence < PRESENCE_THRESHOLD)):
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:
if isinstance(connection_drawing_spec, Mapping):
cv2.line(image, idx_to_coordinates[start_idx],
idx_to_coordinates[end_idx],
connection_drawing_spec[connection].color,
connection_drawing_spec[connection].thickness)
else:
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 idx, landmark_px in idx_to_coordinates.items():
if isinstance(landmark_drawing_spec, Mapping):
cv2.circle(img=image, center=(int(landmark_px[0]) - landmark_drawing_spec[idx].circle_radius//2, int(landmark_px[1])), radius=landmark_drawing_spec[idx].circle_radius//2, color=landmark_drawing_spec[idx].color, thickness=landmark_drawing_spec[idx].thickness)
cv2.circle(img=image, center=(int(landmark_px[0]) + landmark_drawing_spec[idx].circle_radius//2, int(landmark_px[1])), radius=landmark_drawing_spec[idx].circle_radius//2, color=landmark_drawing_spec[idx].color, thickness=landmark_drawing_spec[idx].thickness)
# Triangle (bottom of heart)
pts = np.array([
[int(landmark_px[0]) - landmark_drawing_spec[idx].circle_radius, int(landmark_px[1])],
[int(landmark_px[0]) + landmark_drawing_spec[idx].circle_radius, int(landmark_px[1])],
[int(landmark_px[0]), landmark_px[1] + landmark_drawing_spec[idx].circle_radius*2]
], np.int32).reshape((-1, 1, 2))
cv2.fillPoly(image, [pts], landmark_drawing_spec[idx].color)
else:
# Two circles (top lobes of heart)
cv2.circle(img=image, center=(int(landmark_px[0]) - landmark_drawing_spec.circle_radius//2, int(landmark_px[1])), radius=landmark_drawing_spec.circle_radius//2, color=landmark_drawing_spec.color, thickness=landmark_drawing_spec.thickness)
cv2.circle(img=image, center=(int(landmark_px[0]) + landmark_drawing_spec.circle_radius//2, int(landmark_px[1])), radius=landmark_drawing_spec.circle_radius//2, color=landmark_drawing_spec.color, thickness=landmark_drawing_spec.thickness)
# Triangle (bottom of heart)
pts = np.array([
[int(landmark_px[0]) - landmark_drawing_spec.circle_radius, int(landmark_px[1])],
[int(landmark_px[0]) + landmark_drawing_spec.circle_radius, int(landmark_px[1])],
[int(landmark_px[0]), landmark_px[1] + landmark_drawing_spec.circle_radius*2]
], np.int32).reshape((-1, 1, 2))
cv2.fillPoly(image, [pts], landmark_drawing_spec.color)
#cv2.circle(image, landmark_px, landmark_drawing_spec.circle_radius,
#landmark_drawing_spec.color, landmark_drawing_spec.thickness) |