Update overlay.py

overlay.py

@@ -1,15 +1,13 @@
-# overlay.py
-
 import numpy as np
 from PIL import Image
 import mediapipe as mp
-from baldhead import inference as bald_inference
-from segmentation import remove_face_using_segmentation
 
-# MediaPipe Face Detection for bbox
-mp_fd = mp.solutions.face_detection.FaceDetection(
-    model_selection=1, min_detection_confidence=0.5
-)
+from baldhead import inference       # cạo tóc background
+from segmentation import extract_hair
+
+# MediaPipe Face Detection
+mp_fd = mp.solutions.face_detection.FaceDetection(model_selection=1,
+                                                  min_detection_confidence=0.5)
 
 def get_face_bbox(img: Image.Image) -> tuple[int,int,int,int] | None:
     arr = np.array(img.convert("RGB"))
@@ -29,8 +27,8 @@ def compute_scale(w_bg, h_bg, w_src, h_src) -> float:
 
 def compute_offset(bbox_bg, bbox_src, scale) -> tuple[int,int]:
     x1, y1, x2, y2 = bbox_bg
-    bg_cx = x1 + (x2 - x1) // 2
-    bg_cy = y1 + (y2 - y1) // 2
+    bg_cx = x1 + (x2 - x1)//2
+    bg_cy = y1 + (y2 - y1)//2
     sx1, sy1, sx2, sy2 = bbox_src
     src_cx = int((sx1 + (sx2 - sx1)//2) * scale)
     src_cy = int((sy1 + (sy2 - sy1)//2) * scale)
@@ -47,8 +45,8 @@ def paste_with_alpha(bg: np.ndarray, src: np.ndarray, offset: tuple[int,int]) ->
         return Image.fromarray(res)
     cs = src[y1-y:y2-y, x1-x:x2-x]
     cd = res[y1:y2, x1:x2]
-    mask = cs[...,3]>0
-    if cd.shape[2]==3:
+    mask = cs[...,3] > 0
+    if cd.shape[2] == 3:
         cd[mask] = cs[mask][...,:3]
     else:
         cd[mask] = cs[mask]
@@ -56,7 +54,7 @@ def paste_with_alpha(bg: np.ndarray, src: np.ndarray, offset: tuple[int,int]) ->
     return Image.fromarray(res)
 
 def overlay_source(background: Image.Image, source: Image.Image):
-    # 1) Detect face bboxes
+    # 1) detect bboxes
     bbox_bg  = get_face_bbox(background)
     bbox_src = get_face_bbox(source)
     if bbox_bg is None:
@@ -64,30 +62,28 @@ def overlay_source(background: Image.Image, source: Image.Image):
     if bbox_src is None:
         return None, "❌ No face in source."
 
-    # 2) Scale & resize source
+    # 2) compute scale & resize source
     w_bg, h_bg = bbox_bg[2]-bbox_bg[0], bbox_bg[3]-bbox_bg[1]
     w_src, h_src = bbox_src[2]-bbox_src[0], bbox_src[3]-bbox_src[1]
     scale = compute_scale(w_bg, h_bg, w_src, h_src)
-    # src_scaled = source.resize(
-    #     (int(source.width*scale), int(source.height*scale)),
-    #     Image.Resampling.LANCZOS
-    # )
-    src_scaled = source.resize((int(source.width*scale), int(source.height*scale)))
+    src_scaled = source.resize(
+        (int(source.width*scale), int(source.height*scale)),
+        Image.Resampling.LANCZOS
+    )
 
-    # 3) Compute offset
+    # 3) compute offset
     offset = compute_offset(bbox_bg, bbox_src, scale)
 
-    # 4) BaldHead background using baldhead.py
-    bg_bald = bald_inference(background)
+    # 4) baldhead background
+    bg_bald = inference(background)
 
-    # 5) Extract hair-only from scaled source
-    hair_only = remove_face_using_segmentation(src_scaled)
+    # 5) extract hair-only from source
+    hair_only = extract_hair(src_scaled)
 
-    # 6) Paste hair onto bald background
+    # 6) paste onto bald background
     result = paste_with_alpha(
         np.array(bg_bald.convert("RGBA")),
         np.array(hair_only),
         offset
     )
-
     return result, f"✅ Scale={scale:.2f}, Offset={offset}"