File size: 2,341 Bytes
5029d29
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
"""Draw the holistic skeleton (face mesh, pose, both hands) onto a frame."""

from __future__ import annotations

import numpy as np

from .landmarks import HolisticTracker, LandmarkFrame


def draw_skeleton(
    bgr_frame: np.ndarray,
    frame: LandmarkFrame,
    tracker: HolisticTracker,
    *,
    draw_face: bool = True,
    draw_pose: bool = True,
    draw_hands: bool = True,
) -> np.ndarray:
    """Overlay the detected skeleton on a BGR frame, in place, and return it."""
    if frame.raw is None:
        return bgr_frame

    mp_h = tracker.mp_holistic
    mp_d = tracker.mp_drawing
    mp_s = tracker.mp_styles
    result = frame.raw

    if draw_face and result.face_landmarks is not None:
        # Tesselation gives the full mesh; contours highlight eyes + lips.
        mp_d.draw_landmarks(
            bgr_frame,
            result.face_landmarks,
            mp_h.FACEMESH_TESSELATION,
            landmark_drawing_spec=None,
            connection_drawing_spec=mp_s.get_default_face_mesh_tesselation_style(),
        )
        mp_d.draw_landmarks(
            bgr_frame,
            result.face_landmarks,
            mp_h.FACEMESH_CONTOURS,
            landmark_drawing_spec=None,
            connection_drawing_spec=mp_s.get_default_face_mesh_contours_style(),
        )

    if draw_pose and result.pose_landmarks is not None:
        mp_d.draw_landmarks(
            bgr_frame,
            result.pose_landmarks,
            mp_h.POSE_CONNECTIONS,
            landmark_drawing_spec=mp_s.get_default_pose_landmarks_style(),
        )

    if draw_hands:
        if result.left_hand_landmarks is not None:
            mp_d.draw_landmarks(
                bgr_frame,
                result.left_hand_landmarks,
                mp_h.HAND_CONNECTIONS,
                landmark_drawing_spec=mp_s.get_default_hand_landmarks_style(),
                connection_drawing_spec=mp_s.get_default_hand_connections_style(),
            )
        if result.right_hand_landmarks is not None:
            mp_d.draw_landmarks(
                bgr_frame,
                result.right_hand_landmarks,
                mp_h.HAND_CONNECTIONS,
                landmark_drawing_spec=mp_s.get_default_hand_landmarks_style(),
                connection_drawing_spec=mp_s.get_default_hand_connections_style(),
            )

    return bgr_frame