Create segmentation.py

segmentation.py

@@ -0,0 +1,69 @@
+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))}