Spaces:
Running
on
Zero
Running
on
Zero
| 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]): | |
| # 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 | |