coreml-community/ControlNet-v1-1-Annotators-cpu

1from typing import Mapping2 3import mediapipe as mp4import numpy5 6 7mp_drawing = mp.solutions.drawing_utils8mp_drawing_styles = mp.solutions.drawing_styles9mp_face_detection = mp.solutions.face_detection  # Only for counting faces.10mp_face_mesh = mp.solutions.face_mesh11mp_face_connections = mp.solutions.face_mesh_connections.FACEMESH_TESSELATION12mp_hand_connections = mp.solutions.hands_connections.HAND_CONNECTIONS13mp_body_connections = mp.solutions.pose_connections.POSE_CONNECTIONS14 15DrawingSpec = mp.solutions.drawing_styles.DrawingSpec16PoseLandmark = mp.solutions.drawing_styles.PoseLandmark17 18min_face_size_pixels: int = 6419f_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        img_rgb,89        max_faces: int,90        min_confidence: float91):92    """93    Find up to 'max_faces' inside the provided input image.94    If min_face_size_pixels is provided and nonzero it will be used to filter faces that occupy less than this many95    pixels in the image.96    """97    with mp_face_mesh.FaceMesh(98            static_image_mode=True,99            max_num_faces=max_faces,100            refine_landmarks=True,101            min_detection_confidence=min_confidence,102    ) as facemesh:103        img_height, img_width, img_channels = img_rgb.shape104        assert(img_channels == 3)105 106        results = facemesh.process(img_rgb).multi_face_landmarks107 108        if results is None:109            print("No faces detected in controlnet image for Mediapipe face annotator.")110            return numpy.zeros_like(img_rgb)111 112        # Filter faces that are too small113        filtered_landmarks = []114        for lm in results:115            landmarks = lm.landmark116            face_rect = [117                landmarks[0].x,118                landmarks[0].y,119                landmarks[0].x,120                landmarks[0].y,121            ]  # Left, up, right, down.122            for i in range(len(landmarks)):123                face_rect[0] = min(face_rect[0], landmarks[i].x)124                face_rect[1] = min(face_rect[1], landmarks[i].y)125                face_rect[2] = max(face_rect[2], landmarks[i].x)126                face_rect[3] = max(face_rect[3], landmarks[i].y)127            if min_face_size_pixels > 0:128                face_width = abs(face_rect[2] - face_rect[0])129                face_height = abs(face_rect[3] - face_rect[1])130                face_width_pixels = face_width * img_width131                face_height_pixels = face_height * img_height132                face_size = min(face_width_pixels, face_height_pixels)133                if face_size >= min_face_size_pixels:134                    filtered_landmarks.append(lm)135            else:136                filtered_landmarks.append(lm)137 138        # Annotations are drawn in BGR for some reason, but we don't need to flip a zero-filled image at the start.139        empty = numpy.zeros_like(img_rgb)140 141        # Draw detected faces:142        for face_landmarks in filtered_landmarks:143            mp_drawing.draw_landmarks(144                empty,145                face_landmarks,146                connections=face_connection_spec.keys(),147                landmark_drawing_spec=None,148                connection_drawing_spec=face_connection_spec149            )150            draw_pupils(empty, face_landmarks, iris_landmark_spec, 2)151 152        # Flip BGR back to RGB.153        empty = reverse_channels(empty).copy()154 155        return empty156