Sanjayraju30 commited on
Commit
ffdd09f
·
verified ·
1 Parent(s): 96dad42

Create segmentation.py

Browse files
Files changed (1) hide show
  1. segmentation.py +30 -0
segmentation.py ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import cv2
3
+
4
+ try:
5
+ import mediapipe as mp
6
+ _HAS_MP = True
7
+ except Exception:
8
+ _HAS_MP = False
9
+
10
+ def segment_image(img_np):
11
+ """Returns binary mask for hair/head."""
12
+ h, w = img_np.shape[:2]
13
+ mask = np.zeros((h, w), dtype=np.uint8)
14
+ center = (w // 2, int(h * 0.38))
15
+ axes = (int(w * 0.28), int(h * 0.33))
16
+ cv2.ellipse(mask, center, axes, 0, 0, 360, 255, -1)
17
+ return mask
18
+
19
+ def estimate_landmarks(img_np):
20
+ if not _HAS_MP:
21
+ return None
22
+ mp_face = mp.solutions.face_mesh
23
+ with mp_face.FaceMesh(static_image_mode=True, max_num_faces=1) as fm:
24
+ results = fm.process(cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB))
25
+ if not results.multi_face_landmarks:
26
+ return None
27
+ lm = results.multi_face_landmarks[0]
28
+ xs = [p.x for p in lm.landmark[:10]]
29
+ ys = [p.y for p in lm.landmark[:10]]
30
+ return {"forehead_anchor": (float(np.mean(xs)), float(np.mean(ys)))}