segmentation.py

from transformers import SegformerImageProcessor, AutoModelForSemanticSegmentation
from PIL import Image
import torch
import torch.nn.functional as F
import numpy as np
import mediapipe as mp
import cv2
import os
import warnings

# Suppress MediaPipe warnings
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
warnings.filterwarnings('ignore')

# Suppress MediaPipe logs
import logging
logging.getLogger('mediapipe').setLevel(logging.ERROR)

# Load model
processor = SegformerImageProcessor.from_pretrained("VanNguyen1214/get_face_and_hair")
model = AutoModelForSemanticSegmentation.from_pretrained("VanNguyen1214/get_face_and_hair")

def get_facemesh_mask(image):
    image_np = np.array(image)
    height, width, _ = image_np.shape
    face_mask = np.zeros((height, width), dtype=np.uint8)
    mp_face_mesh = mp.solutions.face_mesh
    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(image_np)
        if results.multi_face_landmarks:
            for face_landmarks in results.multi_face_landmarks:
                points = []
                for lm in face_landmarks.landmark:
                    x, y = int(lm.x * width), int(lm.y * height)
                    points.append([x, y])
                points = np.array(points, np.int32)
                # FaceMesh polygon (bao mặt, trán, không lấy cổ)
                if len(points) > 0:
                    hull = cv2.convexHull(points)
                    cv2.fillConvexPoly(face_mask, hull, 1)
    return face_mask

def expand_forehead_mask(face_mask, expand_percent=0.2):
    ys, xs = np.where(face_mask > 0)
    if len(ys) == 0:
        return face_mask  # không tìm thấy mặt
    min_y, max_y = ys.min(), ys.max()
    height = max_y - min_y
    expand = int(height * expand_percent)
    expanded_min_y = max(min_y - expand, 0)
    expanded_mask = np.zeros_like(face_mask)
    
    # Đảm bảo không lỗi khi vùng mở rộng vượt ngoài ảnh
    src_start = min_y
    src_end = max_y
    dst_start = expanded_min_y
    dst_end = expanded_min_y + (src_end - src_start)
    if dst_end > face_mask.shape[0]:
        overlap = dst_end - face_mask.shape[0]
        dst_end = face_mask.shape[0]
        src_end -= overlap

    expanded_mask[dst_start:dst_end, :] = face_mask[src_start:src_end, :]
    return expanded_mask

def extract_hair(image: Image.Image) -> Image.Image:
    """
    Return an RGBA image where hair pixels have alpha=255 and
    all other pixels have alpha=0.
    """
    rgb = image.convert("RGB")
    arr = np.array(rgb)
    h, w = arr.shape[:2]

    # Segment hair
    inputs = processor(images=rgb, return_tensors="pt")
    with torch.no_grad():
        logits = model(**inputs).logits.cpu()
    up = F.interpolate(logits, size=(h, w), mode="bilinear", align_corners=False)
    seg = up.argmax(dim=1)[0].numpy()
    hair_mask = (seg == 2).astype(np.uint8)

    # Build RGBA
    alpha = (hair_mask * 255).astype(np.uint8)
    rgba  = np.dstack([arr, alpha])
    return Image.fromarray(rgba)

def get_face(image):
    image = image.convert("RGB")
    # SegFormer hair mask
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits.cpu()
    upsampled_logits = F.interpolate(
        logits,
        size=image.size[::-1],
        mode="bilinear",
        align_corners=False,
    )
    pred_seg = upsampled_logits.argmax(dim=1)[0].numpy()
    hair_mask = (pred_seg == 2).astype(np.uint8)  # tóc

    # Face mesh mask (bao trọn mặt, trán, không cổ)
    face_mesh_mask = get_facemesh_mask(image)
    # Expand lên trên 20% chiều cao mặt (ăn gian trán)
    expanded_face_mask = expand_forehead_mask(face_mesh_mask, expand_percent=0.2)

    # Vùng trán mở rộng chỉ lấy phần không trùng vùng mặt gốc và không trùng tóc
    expanded_only_forehead = cv2.bitwise_and(expanded_face_mask, 1 - face_mesh_mask)
    expanded_only_forehead = cv2.bitwise_and(expanded_only_forehead, 1 - hair_mask)

    # Kết hợp: tóc + mặt mediapipe (gốc) + vùng trán mở rộng (phía trên mặt gốc, không trùng tóc, không trùng mặt gốc)
    combined_mask = ((face_mesh_mask + expanded_only_forehead) > 0).astype(np.uint8)

    # Làm mượt mask
    combined_mask = cv2.GaussianBlur(combined_mask.astype(np.float32), (3, 3), 0)
    combined_mask = (combined_mask > 0.5).astype(np.uint8)

    np_image = np.array(image)
    alpha = (combined_mask * 255).astype(np.uint8)
    rgba_image = np.dstack([np_image, alpha])
    return Image.fromarray(rgba_image)

def extract_hair_face_full_forehead(image):
    image = image.convert("RGB")
    # SegFormer hair mask
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits.cpu()
    upsampled_logits = F.interpolate(
        logits,
        size=image.size[::-1],
        mode="bilinear",
        align_corners=False,
    )
    pred_seg = upsampled_logits.argmax(dim=1)[0].numpy()
    hair_mask = (pred_seg == 2).astype(np.uint8)  # tóc

    # Face mesh mask (bao trọn mặt, trán, không cổ)
    face_mesh_mask = get_facemesh_mask(image)
    # Expand lên trên 20% chiều cao mặt (ăn gian trán)
    expanded_face_mask = expand_forehead_mask(face_mesh_mask, expand_percent=0.2)

    # Vùng trán mở rộng chỉ lấy phần không trùng vùng mặt gốc và không trùng tóc
    expanded_only_forehead = cv2.bitwise_and(expanded_face_mask, 1 - face_mesh_mask)
    expanded_only_forehead = cv2.bitwise_and(expanded_only_forehead, 1 - hair_mask)

    # Kết hợp: tóc + mặt mediapipe (gốc) + vùng trán mở rộng (phía trên mặt gốc, không trùng tóc, không trùng mặt gốc)
    combined_mask = ((hair_mask + face_mesh_mask + expanded_only_forehead) > 0).astype(np.uint8)

    # Làm mượt mask
    combined_mask = cv2.GaussianBlur(combined_mask.astype(np.float32), (3, 3), 0)
    combined_mask = (combined_mask > 0.5).astype(np.uint8)

    np_image = np.array(image)
    alpha = (combined_mask * 255).astype(np.uint8)
    rgba_image = np.dstack([np_image, alpha])
    return Image.fromarray(rgba_image)