segmentation.py

import numpy as np
import cv2
try:
    import mediapipe as mp
    _HAS_MP = True
except Exception:
    _HAS_MP = False

def segment_image(img_np):
    """
    Returns a binary mask (uint8 0/255) for hair/head region using MediaPipe if available,
    otherwise falls back to naive ellipse.
    """
    if not _HAS_MP:
        # Fallback to naive ellipse
        h, w = img_np.shape[:2]
        mask = np.zeros((h, w), dtype=np.uint8)
        center = (w // 2, int(h * 0.38))
        axes = (int(w * 0.28), int(h * 0.33))
        cv2.ellipse(mask, center, axes, 0, 0, 360, 255, -1)
        return mask

    # Use MediaPipe FaceMesh for better head segmentation
    mp_face = mp.solutions.face_mesh
    with mp_face.FaceMesh(static_image_mode=True, max_num_faces=1) as face_mesh:
        results = face_mesh.process(cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB))
        if not results.multi_face_landmarks:
            return np.zeros(img_np.shape[:2], dtype=np.uint8)  # Empty mask if no face

        landmarks = results.multi_face_landmarks[0]
        h, w = img_np.shape[:2]
        mask = np.zeros((h, w), dtype=np.uint8)

        # Define head contour points (approximate hair region using face landmarks)
        head_points = [
            10,  # Top forehead
            109, 67, 103, 54, 21, 162, 127, 234, 93, 132, 215,  # Left side
            338, 297, 332, 284, 251, 389, 356, 454, 323, 361, 288,  # Right side
            152  # Chin
        ]
        points = np.array([(int(landmarks.landmark[i].x * w), int(landmarks.landmark[i].y * h)) for i in head_points])

        # Expand slightly for hair
        hull = cv2.convexHull(points)
        cv2.fillConvexPoly(mask, hull, 255)

        # Dilate to cover more hair area
        kernel = np.ones((15, 15), np.uint8)
        mask = cv2.dilate(mask, kernel, iterations=2)

        return mask

def estimate_landmarks(img_np):
    """
    Return dict with key landmarks for alignment (e.g., forehead anchor).
    """
    if not _HAS_MP:
        return None
    mp_face = mp.solutions.face_mesh
    with mp_face.FaceMesh(static_image_mode=True, max_num_faces=1) as face_mesh:
        results = face_mesh.process(cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB))
        if not results.multi_face_landmarks:
            return None
        lm = results.multi_face_landmarks[0].landmark
        # Forehead anchor (average of top landmarks)
        xs = [p.x for p in lm[:10]]  # Normalized [0,1]
        ys = [p.y for p in lm[:10]]
        h, w = img_np.shape[:2]
        return {"forehead_anchor": (int(np.mean(xs) * w), int(np.mean(ys) * h))}