src/segmentation_utils.py

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)