52 lines
3.4 KiB
Python
52 lines
3.4 KiB
Python
import cv2
|
|
import numpy as np
|
|
from typing import Optional, Sequence, Tuple
|
|
|
|
|
|
__all__ = ['get_landmark_connectivity', 'plot_landmarks']
|
|
|
|
|
|
def get_landmark_connectivity(num_landmarks: int) -> Optional[Sequence[Tuple[int, int]]]:
|
|
if num_landmarks == 68:
|
|
return ((0, 1), (1, 2), (2, 3), (3, 4), (4, 5), (5, 6), (6, 7), (7, 8), (8, 9), (9, 10), (10, 11), (11, 12),
|
|
(12, 13), (13, 14), (14, 15), (15, 16), (17, 18), (18, 19), (19, 20), (20, 21), (22, 23), (23, 24),
|
|
(24, 25), (25, 26), (27, 28), (28, 29), (29, 30), (30, 33), (31, 32), (32, 33), (33, 34), (34, 35),
|
|
(36, 37), (37, 38), (38, 39), (40, 41), (41, 36), (42, 43), (43, 44), (44, 45), (45, 46), (46, 47),
|
|
(47, 42), (48, 49), (49, 50), (50, 51), (51, 52), (52, 53), (53, 54), (54, 55), (55, 56), (56, 57),
|
|
(57, 58), (58, 59), (59, 48), (60, 61), (61, 62), (62, 63), (63, 64), (64, 65), (65, 66), (66, 67),
|
|
(67, 60), (39, 40))
|
|
elif num_landmarks == 100:
|
|
return ((0, 1), (1, 2), (2, 3), (3, 4), (4, 5), (5, 6), (6, 7), (7, 8), (8, 9), (9, 10), (10, 11), (11, 12),
|
|
(12, 13), (13, 14), (14, 15), (15, 16), (17, 18), (18, 19), (19, 20), (20, 21), (22, 23), (23, 24),
|
|
(24, 25), (25, 26), (68, 69), (69, 70), (70, 71), (72, 73), (73, 74), (74, 75), (36, 76), (76, 37),
|
|
(37, 77), (77, 38), (38, 78), (78, 39), (39, 40), (40, 79), (79, 41), (41, 36), (42, 80), (80, 43),
|
|
(43, 81), (81, 44), (44, 82), (82, 45), (45, 46), (46, 83), (83, 47), (47, 42), (27, 28), (28, 29),
|
|
(29, 30), (30, 33), (31, 32), (32, 33), (33, 34), (34, 35), (84, 85), (86, 87), (48, 49), (49, 88),
|
|
(88, 50), (50, 51), (51, 52), (52, 89), (89, 53), (53, 54), (54, 55), (55, 90), (90, 56), (56, 57),
|
|
(57, 58), (58, 91), (91, 59), (59, 48), (60, 92), (92, 93), (93, 61), (61, 62), (62, 63), (63, 94),
|
|
(94, 95), (95, 64), (64, 96), (96, 97), (97, 65), (65, 66), (66, 67), (67, 98), (98, 99), (99, 60),
|
|
(17, 68), (21, 71), (22, 72), (26, 75))
|
|
else:
|
|
return None
|
|
|
|
|
|
def plot_landmarks(image: np.ndarray, landmarks: np.ndarray, landmark_scores: Optional[Sequence[float]] = None,
|
|
threshold: float = 0.2, line_colour: Tuple[int, int, int] = (0, 255, 0),
|
|
pts_colour: Tuple[int, int, int] = (0, 0, 255), line_thickness: int = 1, pts_radius: int = 1,
|
|
landmark_connectivity: Optional[Sequence[Tuple[int, int]]] = None) -> None:
|
|
num_landmarks = len(landmarks)
|
|
if landmark_scores is None:
|
|
landmark_scores = np.full((num_landmarks,), threshold + 1.0, dtype=float)
|
|
if landmark_connectivity is None:
|
|
landmark_connectivity = get_landmark_connectivity(len(landmarks))
|
|
if landmark_connectivity is not None:
|
|
for (idx1, idx2) in landmark_connectivity:
|
|
if (idx1 < num_landmarks and idx2 < num_landmarks and
|
|
landmark_scores[idx1] >= threshold and landmark_scores[idx2] >= threshold):
|
|
cv2.line(image, tuple(landmarks[idx1].astype(int).tolist()),
|
|
tuple(landmarks[idx2].astype(int).tolist()),
|
|
color=line_colour, thickness=line_thickness, lineType=cv2.LINE_AA)
|
|
for landmark, score in zip(landmarks, landmark_scores):
|
|
if score >= threshold:
|
|
cv2.circle(image, tuple(landmark.astype(int).tolist()), pts_radius, pts_colour, -1)
|