from mediapipe.tasks import python from mediapipe.tasks.python import vision import mediapipe as mp import cv2 import numpy as np import warnings import os import logging # Suppress INFO and WARNING logs from MediaPipe logging.getLogger("mediapipe").setLevel(logging.ERROR) # Suppress INFO and WARNING logs os.environ["GLOG_minloglevel"] = "2" # 2 means only ERROR and FATAL logs os.environ["GLOG_logtostderr"] = "1" # Initialize mediapipe solutions mp_face_detection = mp.solutions.face_detection # type: ignore mp_face_mesh = mp.solutions.face_mesh # type: ignore def detect_faces_and_landmarks(image: np.ndarray): """ Detect faces and landmarks using MediaPipe Face Detection. :param image: Input image as a numpy array. :return: List of dictionaries with face and landmark information. """ with mp_face_detection.FaceDetection( model_selection=1, min_detection_confidence=0.5 ) as face_detection: results = face_detection.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)) face_data = [] if results.detections: for detection in results.detections: bboxC = detection.location_data.relative_bounding_box h, w, c = image.shape bbox = ( int(bboxC.xmin * w), int(bboxC.ymin * h), int(bboxC.width * w), int(bboxC.height * h), ) landmarks = detection.location_data.relative_keypoints face_data.append({"bbox": bbox, "landmarks": landmarks}) return face_data def mediapipe_selfie_segmentor( image: np.ndarray, segment: list = ["face_skin", "body_skin", "hair"] ): """ Segment image using MediaPipe Multi-Class Selfie Segmentation. :param image: Input image as a numpy array. :param segment: List of segments to extract. :return: Dictionary of segmentation masks. """ # Create the options that will be used for ImageSegmenter base_options = python.BaseOptions( model_asset_path="model_weights/selfie_multiclass_256x256.tflite" ) options = vision.ImageSegmenterOptions( base_options=base_options, output_category_mask=True, output_confidence_masks=True, ) with vision.ImageSegmenter.create_from_options(options) as segmenter: # Create the MediaPipe image file that will be segmented mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=image) # Retrieve the masks for the segmented image segmentation_result = segmenter.segment(mp_image) category_mask = segmentation_result.category_mask.numpy_view() h, w = category_mask.shape masks = { "face_skin_mask": np.zeros((h, w), dtype=np.uint8), "hair_mask": np.zeros((h, w), dtype=np.uint8), "body_skin_mask": np.zeros((h, w), dtype=np.uint8), } # Define class labels based on MediaPipe segmentation (example, may need adjustment) face_skin_class = 3 hair_class = 1 body_skin_class = 2 masks["face_skin_mask"][category_mask == face_skin_class] = 255 masks["hair_mask"][category_mask == hair_class] = 255 masks["body_skin_mask"][category_mask == body_skin_class] = 255 return masks def detect_face_landmarks(image: np.ndarray): """ Detect face landmarks using MediaPipe Face Mesh. :param image: Input image as a numpy array. :return: Dictionary with landmarks for iris, lips, eyebrows, and eyes. """ with mp_face_mesh.FaceMesh( static_image_mode=True, max_num_faces=1, refine_landmarks=True, min_detection_confidence=0.5, ) as face_mesh: results = face_mesh.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)) face_landmarks = { "left_iris": [], "right_iris": [], "lips": [], "left_eyebrow": [], "right_eyebrow": [], "left_eye": [], "right_eye": [], } if results.multi_face_landmarks: for face_landmarks_data in results.multi_face_landmarks: # Left iris landmarks for i in range(468, 473): # Left iris landmarks landmark = face_landmarks_data.landmark[i] face_landmarks["left_iris"].append((landmark.x, landmark.y)) # Right iris landmarks for i in range(473, 478): # Right iris landmarks landmark = face_landmarks_data.landmark[i] face_landmarks["right_iris"].append((landmark.x, landmark.y)) # Outer lips landmarks for i in [ 61, 146, 91, 181, 84, 17, 314, 405, 321, 375, 291, 0, 409, 270, 269, 267, 37, 39, 40, 185, ]: landmark = face_landmarks_data.landmark[i] face_landmarks["lips"].append((landmark.x, landmark.y)) # Left eyebrow landmarks for i in [70, 63, 105, 66, 107]: landmark = face_landmarks_data.landmark[i] face_landmarks["left_eyebrow"].append((landmark.x, landmark.y)) # Right eyebrow landmarks for i in [336, 296, 334, 293, 300]: landmark = face_landmarks_data.landmark[i] face_landmarks["right_eyebrow"].append((landmark.x, landmark.y)) # Left eye landmarks for i in [ 33, 246, 161, 160, 159, 158, 157, 173, 133, 155, 154, 153, 145, 144, 163, 7, ]: landmark = face_landmarks_data.landmark[i] face_landmarks["left_eye"].append((landmark.x, landmark.y)) # Right eye landmarks for i in [ 463, 398, 384, 385, 386, 387, 388, 466, 263, 249, 390, 373, 374, 380, 381, 382, ]: landmark = face_landmarks_data.landmark[i] face_landmarks["right_eye"].append((landmark.x, landmark.y)) return face_landmarks def create_feature_masks(image: np.ndarray, landmarks: dict): """ Create individual masks for facial features based on landmarks. :param image: Input image as a numpy array. :param landmarks: Dictionary with landmarks for iris, lips, eyebrows, and eyes. :return: Dictionary with masks for each facial feature. """ h, w = image.shape[:2] masks = { "lips_mask": np.zeros((h, w), dtype=np.uint8), "left_eyebrow_mask": np.zeros((h, w), dtype=np.uint8), "right_eyebrow_mask": np.zeros((h, w), dtype=np.uint8), "left_eye_mask": np.zeros((h, w), dtype=np.uint8), "right_eye_mask": np.zeros((h, w), dtype=np.uint8), "left_iris_mask": np.zeros((h, w), dtype=np.uint8), "right_iris_mask": np.zeros((h, w), dtype=np.uint8), } # Define the order of the points to form polygons correctly lips_order = [ 61, 146, 91, 181, 84, 17, 314, 405, 321, 375, 291, 0, 409, 270, 269, 267, 37, 39, 40, 185, ] left_eyebrow_order = [70, 63, 105, 66, 107] right_eyebrow_order = [336, 296, 334, 293, 300] left_eye_order = [ 33, 246, 161, 160, 159, 158, 157, 173, 133, 155, 154, 153, 145, 144, 163, 7, ] right_eye_order = [ 463, 398, 384, 385, 386, 387, 388, 466, 263, 249, 390, 373, 374, 380, 381, 382, ] left_iris_order = [468, 469, 470, 471, 472] right_iris_order = [473, 474, 475, 476, 477] orders = { "lips": lips_order, "left_eyebrow": left_eyebrow_order, "right_eyebrow": right_eyebrow_order, "left_eye": left_eye_order, "right_eye": right_eye_order, "left_iris": left_iris_order, "right_iris": right_iris_order, } for feature, order in orders.items(): points = [] for i in range(len(order)): try: point = ( int(landmarks[feature][i][0] * w), int(landmarks[feature][i][1] * h), ) points.append(point) except KeyError: warnings.warn( f"Feature '{feature}' at index {i} is not present in landmarks. Skipping this point." ) except IndexError: warnings.warn( f"Index {i} is out of range for feature '{feature}'. Skipping this point." ) points = np.array(points, dtype=np.int32) if len(points) > 0: cv2.fillPoly(masks[f"{feature}_mask"], [points], 255) return masks if __name__ == "__main__": # Test the face detection and segmentation image = cv2.imread("inputs/vanika.png") face_data = detect_faces_and_landmarks(image) print(face_data) masks = mediapipe_selfie_segmentor(image) # write it to disk for key, mask in masks.items(): if key == "face_skin_mask": # create feature masks landmarks = detect_face_landmarks(image) feature_masks = create_feature_masks(image, landmarks) # subtract eyes, lips and eyebrows from face skin mask for feature, feature_mask in feature_masks.items(): if "iris_mask" in feature: cv2.imwrite(f"outputs/{feature}.png", feature_mask) mask = cv2.subtract(mask, feature_mask) cv2.imwrite(f"outputs/{key}.png", mask)