CrucibleAI/ControlNetMediaPipeFace

1from typing import Mapping2 3import mediapipe as mp4import numpy5from PIL import Image6 7 8mp_drawing = mp.solutions.drawing_utils9mp_drawing_styles = mp.solutions.drawing_styles10mp_face_detection = mp.solutions.face_detection  # Only for counting faces.11mp_face_mesh = mp.solutions.face_mesh12mp_face_connections = mp.solutions.face_mesh_connections.FACEMESH_TESSELATION13mp_hand_connections = mp.solutions.hands_connections.HAND_CONNECTIONS14mp_body_connections = mp.solutions.pose_connections.POSE_CONNECTIONS15 16DrawingSpec = mp.solutions.drawing_styles.DrawingSpec17PoseLandmark = mp.solutions.drawing_styles.PoseLandmark18 19f_thick = 220f_rad = 121right_iris_draw = DrawingSpec(color=(10, 200, 250), thickness=f_thick, circle_radius=f_rad)22right_eye_draw = DrawingSpec(color=(10, 200, 180), thickness=f_thick, circle_radius=f_rad)23right_eyebrow_draw = DrawingSpec(color=(10, 220, 180), thickness=f_thick, circle_radius=f_rad)24left_iris_draw = DrawingSpec(color=(250, 200, 10), thickness=f_thick, circle_radius=f_rad)25left_eye_draw = DrawingSpec(color=(180, 200, 10), thickness=f_thick, circle_radius=f_rad)26left_eyebrow_draw = DrawingSpec(color=(180, 220, 10), thickness=f_thick, circle_radius=f_rad)27mouth_draw = DrawingSpec(color=(10, 180, 10), thickness=f_thick, circle_radius=f_rad)28head_draw = DrawingSpec(color=(10, 200, 10), thickness=f_thick, circle_radius=f_rad)29 30# mp_face_mesh.FACEMESH_CONTOURS has all the items we care about.31face_connection_spec = {}32for edge in mp_face_mesh.FACEMESH_FACE_OVAL:33    face_connection_spec[edge] = head_draw34for edge in mp_face_mesh.FACEMESH_LEFT_EYE:35    face_connection_spec[edge] = left_eye_draw36for edge in mp_face_mesh.FACEMESH_LEFT_EYEBROW:37    face_connection_spec[edge] = left_eyebrow_draw38# for edge in mp_face_mesh.FACEMESH_LEFT_IRIS:39#    face_connection_spec[edge] = left_iris_draw40for edge in mp_face_mesh.FACEMESH_RIGHT_EYE:41    face_connection_spec[edge] = right_eye_draw42for edge in mp_face_mesh.FACEMESH_RIGHT_EYEBROW:43    face_connection_spec[edge] = right_eyebrow_draw44# for edge in mp_face_mesh.FACEMESH_RIGHT_IRIS:45#    face_connection_spec[edge] = right_iris_draw46for edge in mp_face_mesh.FACEMESH_LIPS:47    face_connection_spec[edge] = mouth_draw48iris_landmark_spec = {468: right_iris_draw, 473: left_iris_draw}49 50 51def draw_pupils(image, landmark_list, drawing_spec, halfwidth: int = 2):52    """We have a custom function to draw the pupils because the mp.draw_landmarks method requires a parameter for all53    landmarks.  Until our PR is merged into mediapipe, we need this separate method."""54    if len(image.shape) != 3:55        raise ValueError("Input image must be H,W,C.")56    image_rows, image_cols, image_channels = image.shape57    if image_channels != 3:  # BGR channels58        raise ValueError('Input image must contain three channel bgr data.')59    for idx, landmark in enumerate(landmark_list.landmark):60        if (61                (landmark.HasField('visibility') and landmark.visibility < 0.9) or62                (landmark.HasField('presence') and landmark.presence < 0.5)63        ):64            continue65        if landmark.x >= 1.0 or landmark.x < 0 or landmark.y >= 1.0 or landmark.y < 0:66            continue67        image_x = int(image_cols*landmark.x)68        image_y = int(image_rows*landmark.y)69        draw_color = None70        if isinstance(drawing_spec, Mapping):71            if drawing_spec.get(idx) is None:72                continue73            else:74                draw_color = drawing_spec[idx].color75        elif isinstance(drawing_spec, DrawingSpec):76            draw_color = drawing_spec.color77        image[image_y-halfwidth:image_y+halfwidth, image_x-halfwidth:image_x+halfwidth, :] = draw_color78 79 80def reverse_channels(image):81    """Given a numpy array in RGB form, convert to BGR.  Will also convert from BGR to RGB."""82    # im[:,:,::-1] is a neat hack to convert BGR to RGB by reversing the indexing order.83    # im[:,:,::[2,1,0]] would also work but makes a copy of the data.84    return image[:, :, ::-1]85 86 87def generate_annotation(88        input_image: Image.Image,89        max_faces: int,90        min_face_size_pixels: int = 0,91        return_annotation_data: bool = False92):93    """94    Find up to 'max_faces' inside the provided input image.95    If min_face_size_pixels is provided and nonzero it will be used to filter faces that occupy less than this many96    pixels in the image.97    If return_annotation_data is TRUE (default: false) then in addition to returning the 'detected face' image, three98    additional parameters will be returned: faces before filtering, faces after filtering, and an annotation image.99    The faces_before_filtering return value is the number of faces detected in an image with no filtering.100    faces_after_filtering is the number of faces remaining after filtering small faces.101 102    :return:103      If 'return_annotation_data==True', returns (numpy array, numpy array, int, int).104      If 'return_annotation_data==False' (default), returns a numpy array.105    """106    with mp_face_mesh.FaceMesh(107            static_image_mode=True,108            max_num_faces=max_faces,109            refine_landmarks=True,110            min_detection_confidence=0.5,111    ) as facemesh:112        img_rgb = numpy.asarray(input_image)113        results = facemesh.process(img_rgb).multi_face_landmarks114 115        faces_found_before_filtering = len(results)116 117        # Filter faces that are too small118        filtered_landmarks = []119        for lm in results:120            landmarks = lm.landmark121            face_rect = [122                landmarks[0].x,123                landmarks[0].y,124                landmarks[0].x,125                landmarks[0].y,126            ]  # Left, up, right, down.127            for i in range(len(landmarks)):128                face_rect[0] = min(face_rect[0], landmarks[i].x)129                face_rect[1] = min(face_rect[1], landmarks[i].y)130                face_rect[2] = max(face_rect[2], landmarks[i].x)131                face_rect[3] = max(face_rect[3], landmarks[i].y)132            if min_face_size_pixels > 0:133                face_width = abs(face_rect[2] - face_rect[0])134                face_height = abs(face_rect[3] - face_rect[1])135                face_width_pixels = face_width * input_image.size[0]136                face_height_pixels = face_height * input_image.size[1]137                face_size = min(face_width_pixels, face_height_pixels)138                if face_size >= min_face_size_pixels:139                    filtered_landmarks.append(lm)140            else:141                filtered_landmarks.append(lm)142 143        faces_remaining_after_filtering = len(filtered_landmarks)144 145        # Annotations are drawn in BGR for some reason, but we don't need to flip a zero-filled image at the start.146        empty = numpy.zeros_like(img_rgb)147 148        # Draw detected faces:149        for face_landmarks in filtered_landmarks:150            mp_drawing.draw_landmarks(151                empty,152                face_landmarks,153                connections=face_connection_spec.keys(),154                landmark_drawing_spec=None,155                connection_drawing_spec=face_connection_spec156            )157            draw_pupils(empty, face_landmarks, iris_landmark_spec, 2)158 159        # Flip BGR back to RGB.160        empty = reverse_channels(empty)161 162        # We might have to generate a composite.163        if return_annotation_data:164            # Note that we're copying the input image AND flipping the channels so we can draw on top of it.165            annotated = reverse_channels(numpy.asarray(input_image)).copy()166            for face_landmarks in filtered_landmarks:167                mp_drawing.draw_landmarks(168                    empty,169                    face_landmarks,170                    connections=face_connection_spec.keys(),171                    landmark_drawing_spec=None,172                    connection_drawing_spec=face_connection_spec173                )174                draw_pupils(empty, face_landmarks, iris_landmark_spec, 2)175            annotated = reverse_channels(annotated)176 177        if not return_annotation_data:178            return empty179        else:180            return empty, annotated, faces_found_before_filtering, faces_remaining_after_filtering181