File size: 9,328 Bytes
03a856a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
import cv2
import mediapipe as mp
import numpy as np
from mediapipe.framework.formats import landmark_pb2

class FaceMeshVisualizer:
    def __init__(self, 
            forehead_edge=False, 
            upface_only=False, 
            draw_eye=True,
            draw_head=False,
            draw_iris=True,
            draw_eyebrow=True,
            draw_mouse=True,
            draw_nose=True,
            draw_pupil=True
        ):
        self.mp_drawing = mp.solutions.drawing_utils
        mp_face_mesh = mp.solutions.face_mesh
        self.mp_face_mesh = mp_face_mesh
        self.forehead_edge = forehead_edge

        DrawingSpec = mp.solutions.drawing_styles.DrawingSpec
        f_thick = 2
        f_rad = 1
        right_iris_draw = DrawingSpec(color=(10, 200, 250), thickness=f_thick, circle_radius=f_rad)
        right_eye_draw = DrawingSpec(color=(10, 200, 180), thickness=f_thick, circle_radius=f_rad)
        right_eyebrow_draw = DrawingSpec(color=(10, 220, 180), thickness=f_thick, circle_radius=f_rad)
        left_iris_draw = DrawingSpec(color=(250, 200, 10), thickness=f_thick, circle_radius=f_rad)
        left_eye_draw = DrawingSpec(color=(180, 200, 10), thickness=f_thick, circle_radius=f_rad)
        left_eyebrow_draw = DrawingSpec(color=(180, 220, 10), thickness=f_thick, circle_radius=f_rad)
        head_draw = DrawingSpec(color=(10, 200, 10), thickness=f_thick, circle_radius=f_rad)
        nose_draw = DrawingSpec(color=(200, 200, 200), thickness=f_thick, circle_radius=f_rad)
        
        mouth_draw_obl = DrawingSpec(color=(10, 180, 20), thickness=f_thick, circle_radius=f_rad)
        mouth_draw_obr = DrawingSpec(color=(20, 10, 180), thickness=f_thick, circle_radius=f_rad)
        
        mouth_draw_ibl = DrawingSpec(color=(100, 100, 30), thickness=f_thick, circle_radius=f_rad)
        mouth_draw_ibr = DrawingSpec(color=(100, 150, 50), thickness=f_thick, circle_radius=f_rad)
        
        mouth_draw_otl = DrawingSpec(color=(20, 80, 100), thickness=f_thick, circle_radius=f_rad)
        mouth_draw_otr = DrawingSpec(color=(80, 100, 20), thickness=f_thick, circle_radius=f_rad)
        
        mouth_draw_itl = DrawingSpec(color=(120, 100, 200), thickness=f_thick, circle_radius=f_rad)
        mouth_draw_itr = DrawingSpec(color=(150 ,120, 100), thickness=f_thick, circle_radius=f_rad)
        
        FACEMESH_LIPS_OUTER_BOTTOM_LEFT = [(61,146),(146,91),(91,181),(181,84),(84,17)]
        FACEMESH_LIPS_OUTER_BOTTOM_RIGHT = [(17,314),(314,405),(405,321),(321,375),(375,291)]
        
        FACEMESH_LIPS_INNER_BOTTOM_LEFT = [(78,95),(95,88),(88,178),(178,87),(87,14)]
        FACEMESH_LIPS_INNER_BOTTOM_RIGHT = [(14,317),(317,402),(402,318),(318,324),(324,308)]
        
        FACEMESH_LIPS_OUTER_TOP_LEFT = [(61,185),(185,40),(40,39),(39,37),(37,0)]
        FACEMESH_LIPS_OUTER_TOP_RIGHT = [(0,267),(267,269),(269,270),(270,409),(409,291)]
        
        FACEMESH_LIPS_INNER_TOP_LEFT = [(78,191),(191,80),(80,81),(81,82),(82,13)]
        FACEMESH_LIPS_INNER_TOP_RIGHT = [(13,312),(312,311),(311,310),(310,415),(415,308)]
        
        FACEMESH_CUSTOM_FACE_OVAL = [(176, 149), (150, 136), (356, 454), (58, 132), (152, 148), (361, 288), (251, 389), (132, 93), (389, 356), (400, 377), (136, 172), (377, 152), (323, 361), (172, 58), (454, 323), (365, 379), (379, 378), (148, 176), (93, 234), (397, 365), (149, 150), (288, 397), (234, 127), (378, 400), (127, 162), (162, 21)]

        # mp_face_mesh.FACEMESH_CONTOURS has all the items we care about.
        face_connection_spec = {}

        #from IPython import embed
        #embed()
        if self.forehead_edge:
            for edge in mp_face_mesh.FACEMESH_FACE_OVAL:
                face_connection_spec[edge] = head_draw
        else:
            if draw_head:
                FACEMESH_CUSTOM_FACE_OVAL_sorted = sorted(FACEMESH_CUSTOM_FACE_OVAL)
                if upface_only:
                    for edge in [FACEMESH_CUSTOM_FACE_OVAL_sorted[edge_idx] for edge_idx in [1,2,9,12,13,16,22,25]]:
                        face_connection_spec[edge] = head_draw
                else:
                    for edge in FACEMESH_CUSTOM_FACE_OVAL_sorted:
                        face_connection_spec[edge] = head_draw
        
        if draw_eye:  
            for edge in mp_face_mesh.FACEMESH_LEFT_EYE:
                face_connection_spec[edge] = left_eye_draw
            for edge in mp_face_mesh.FACEMESH_RIGHT_EYE:
                face_connection_spec[edge] = right_eye_draw
        
        if draw_eyebrow:
            for edge in mp_face_mesh.FACEMESH_LEFT_EYEBROW:
                face_connection_spec[edge] = left_eyebrow_draw

            for edge in mp_face_mesh.FACEMESH_RIGHT_EYEBROW:
                face_connection_spec[edge] = right_eyebrow_draw
            
        if draw_iris:
            for edge in mp_face_mesh.FACEMESH_LEFT_IRIS:
                face_connection_spec[edge] = left_iris_draw
            for edge in mp_face_mesh.FACEMESH_RIGHT_IRIS:
                face_connection_spec[edge] = right_iris_draw
            
            #for edge in mp_face_mesh.FACEMESH_RIGHT_EYEBROW:
            #    face_connection_spec[edge] = right_eyebrow_draw
            # for edge in mp_face_mesh.FACEMESH_RIGHT_IRIS:
            #    face_connection_spec[edge] = right_iris_draw

        # for edge in mp_face_mesh.FACEMESH_LIPS:
        #     face_connection_spec[edge] = mouth_draw

        if draw_mouse:
            for edge in FACEMESH_LIPS_OUTER_BOTTOM_LEFT:
                face_connection_spec[edge] = mouth_draw_obl
            for edge in FACEMESH_LIPS_OUTER_BOTTOM_RIGHT:
                face_connection_spec[edge] = mouth_draw_obr
            for edge in FACEMESH_LIPS_INNER_BOTTOM_LEFT:
                face_connection_spec[edge] = mouth_draw_ibl
            for edge in FACEMESH_LIPS_INNER_BOTTOM_RIGHT:
                face_connection_spec[edge] = mouth_draw_ibr
            for edge in FACEMESH_LIPS_OUTER_TOP_LEFT:
                face_connection_spec[edge] = mouth_draw_otl
            for edge in FACEMESH_LIPS_OUTER_TOP_RIGHT:
                face_connection_spec[edge] = mouth_draw_otr
            for edge in FACEMESH_LIPS_INNER_TOP_LEFT:
                face_connection_spec[edge] = mouth_draw_itl
            for edge in FACEMESH_LIPS_INNER_TOP_RIGHT:
                face_connection_spec[edge] = mouth_draw_itr
        
        self.face_connection_spec = face_connection_spec
        
        self.pupil_landmark_spec = {468: right_iris_draw, 473: left_iris_draw}
        self.nose_landmark_spec = {4: nose_draw}

        self.draw_pupil = draw_pupil
        self.draw_nose = draw_nose
    
    def draw_points(self, image, landmark_list, drawing_spec, halfwidth: int = 2):
        """We have a custom function to draw the pupils because the mp.draw_landmarks method requires a parameter for all
        landmarks.  Until our PR is merged into mediapipe, we need this separate method."""
        if len(image.shape) != 3:
            raise ValueError("Input image must be H,W,C.")
        image_rows, image_cols, image_channels = image.shape
        if image_channels != 3:  # BGR channels
            raise ValueError('Input image must contain three channel bgr data.')
        for idx, landmark in enumerate(landmark_list.landmark):
            if idx not in drawing_spec:
                continue
            
            if (
                    (landmark.HasField('visibility') and landmark.visibility < 0.9) or
                    (landmark.HasField('presence') and landmark.presence < 0.5)
            ):
                continue
            if landmark.x >= 1.0 or landmark.x < 0 or landmark.y >= 1.0 or landmark.y < 0:
                continue

            image_x = int(image_cols * landmark.x)
            image_y = int(image_rows * landmark.y)
            
            draw_color = drawing_spec[idx].color
            image[image_y - halfwidth : image_y + halfwidth, image_x - halfwidth : image_x + halfwidth, :] = draw_color
        

    def draw_landmarks(self, image_size, keypoints, normed=False):
        ini_size = [512, 512]
        image = np.zeros([ini_size[1], ini_size[0], 3], dtype=np.uint8)
        if keypoints is not None:
            new_landmarks = landmark_pb2.NormalizedLandmarkList()
            for i in range(keypoints.shape[0]):
                landmark = new_landmarks.landmark.add()
                if normed:
                    landmark.x = keypoints[i, 0]
                    landmark.y = keypoints[i, 1]
                else:
                    landmark.x = keypoints[i, 0] / image_size[0]
                    landmark.y = keypoints[i, 1] / image_size[1]
                landmark.z = 1.0

            self.mp_drawing.draw_landmarks(
                image=image,
                landmark_list=new_landmarks,
                connections=self.face_connection_spec.keys(),
                landmark_drawing_spec=None,
                connection_drawing_spec=self.face_connection_spec
            )

            if self.draw_pupil:
                self.draw_points(image, new_landmarks, self.pupil_landmark_spec, 3)
            
            if self.draw_nose:
                self.draw_points(image, new_landmarks, self.nose_landmark_spec, 3)

        image = cv2.resize(image, (image_size[0], image_size[1]))
        
        return image