diff --git a/adetailer/common.py b/adetailer/common.py index 481adbb..a29bdd5 100644 --- a/adetailer/common.py +++ b/adetailer/common.py @@ -58,6 +58,7 @@ def get_models( "mediapipe_face_full": None, "mediapipe_face_short": None, "mediapipe_face_mesh": None, + "mediapipe_face_mesh_eyes_only": None, } ) diff --git a/adetailer/mediapipe.py b/adetailer/mediapipe.py index 0ec6b16..194496c 100644 --- a/adetailer/mediapipe.py +++ b/adetailer/mediapipe.py @@ -16,6 +16,7 @@ def mediapipe_predict( "mediapipe_face_short": partial(mediapipe_face_detection, 0), "mediapipe_face_full": partial(mediapipe_face_detection, 1), "mediapipe_face_mesh": mediapipe_face_mesh, + "mediapipe_face_mesh_eyes_only": mediapipe_face_mesh_eyes_only, } if model_type in mapping: func = mapping[model_type] @@ -66,9 +67,25 @@ def mediapipe_face_detection( return PredictOutput(bboxes=bboxes, masks=masks, preview=preview) +def get_convexhull(points: np.ndarray) -> list[tuple[int, int]]: + """ + Parameters + ---------- + points: An ndarray of shape (n, 2) containing the 2D points. + + Returns + ------- + list[tuple[int, int]]: Input for the draw.polygon function + """ + from scipy.spatial import ConvexHull + + hull = ConvexHull(points) + vertices = hull.vertices + return list(zip(points[vertices, 0], points[vertices, 1])) + + def mediapipe_face_mesh(image: Image.Image, confidence: float = 0.3) -> PredictOutput: import mediapipe as mp - from scipy.spatial import ConvexHull mp_face_mesh = mp.solutions.face_mesh draw_util = mp.solutions.drawing_utils @@ -98,9 +115,7 @@ def mediapipe_face_mesh(image: Image.Image, confidence: float = 0.3) -> PredictO ) points = np.array([(land.x * w, land.y * h) for land in landmarks.landmark]) - hull = ConvexHull(points) - vertices = hull.vertices - outline = list(zip(points[vertices, 0], points[vertices, 1])) + outline = get_convexhull(points) mask = Image.new("L", image.size, "black") draw = ImageDraw.Draw(mask) @@ -110,3 +125,60 @@ def mediapipe_face_mesh(image: Image.Image, confidence: float = 0.3) -> PredictO bboxes = create_bbox_from_mask(masks, image.size) preview = Image.fromarray(preview) return PredictOutput(bboxes=bboxes, masks=masks, preview=preview) + + +def mediapipe_face_mesh_eyes_only( + image: Image.Image, confidence: float = 0.3 +) -> PredictOutput: + import mediapipe as mp + + mp_face_mesh = mp.solutions.face_mesh + + left_idx = np.array(list(mp_face_mesh.FACEMESH_LEFT_EYE)).flatten() + right_idx = np.array(list(mp_face_mesh.FACEMESH_RIGHT_EYE)).flatten() + + w, h = image.size + + with mp_face_mesh.FaceMesh( + static_image_mode=True, max_num_faces=20, min_detection_confidence=confidence + ) as face_mesh: + arr = np.array(image) + pred = face_mesh.process(arr) + + if pred.multi_face_landmarks is None: + return PredictOutput() + + preview = image.copy() + masks = [] + + for landmarks in pred.multi_face_landmarks: + points = np.array([(land.x * w, land.y * h) for land in landmarks.landmark]) + left_eyes = points[left_idx] + right_eyes = points[right_idx] + left_outline = get_convexhull(left_eyes) + right_outline = get_convexhull(right_eyes) + + mask = Image.new("L", image.size, "black") + draw = ImageDraw.Draw(mask) + for outline in (left_outline, right_outline): + draw.polygon(outline, fill="white") + masks.append(mask) + + bboxes = create_bbox_from_mask(masks, image.size) + preview = draw_preview(preview, bboxes, masks) + return PredictOutput(bboxes=bboxes, masks=masks, preview=preview) + + +def draw_preview( + preview: Image.Image, bboxes: list[list[int]], masks: list[Image.Image] +) -> Image.Image: + for mask in masks: + red = Image.new("RGB", preview.size, "red") + masked = Image.composite(red, preview, mask) + preview = Image.blend(preview, masked, 0.25) + + draw = ImageDraw.Draw(preview) + for bbox in bboxes: + draw.rectangle(bbox, outline="red", width=2) + + return preview