Spaces:
Running
Running
File size: 2,815 Bytes
69fb66d 36c2948 6e97dbf 69fb66d c20459b 6e97dbf 36c2948 6e97dbf c20459b 6e97dbf 1249491 6e97dbf 7f95e91 6e97dbf c20459b 6e97dbf c20459b 6e97dbf c20459b 2c09205 c20459b 36c2948 c20459b 36c2948 c20459b 6e97dbf c20459b 69fb66d c20459b 69fb66d 6e97dbf 69fb66d 390e44b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
import numpy as np
from PIL import Image
import mediapipe as mp
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"))
res = mp_fd.process(arr)
if not res.detections:
return None
d = res.detections[0].location_data.relative_bounding_box
h, w = arr.shape[:2]
x1 = int(d.xmin * w)
y1 = int(d.ymin * h)
x2 = x1 + int(d.width * w)
y2 = y1 + int(d.height * h)
return x1, y1, x2, y2
def compute_scale(w_bg, h_bg, w_src, h_src) -> float:
return ((w_bg / w_src) + (h_bg / h_src)) / 2
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
sx1, sy1, sx2, sy2 = bbox_src
src_cx = int((sx1 + (sx2 - sx1)//2) * scale)
src_cy = int((sy1 + (sy2 - sy1)//2) * scale)
return bg_cx - src_cx, bg_cy - src_cy
def paste_with_alpha(bg: np.ndarray, src: np.ndarray, offset: tuple[int,int]) -> Image.Image:
res = bg.copy()
x, y = offset
h, w = src.shape[:2]
x1, y1 = max(x,0), max(y,0)
x2 = min(x+w, bg.shape[1])
y2 = min(y+h, bg.shape[0])
if x1>=x2 or y1>=y2:
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:
cd[mask] = cs[mask][...,:3]
else:
cd[mask] = cs[mask]
res[y1:y2, x1:x2] = cd
return Image.fromarray(res)
def overlay_source(background: Image.Image, source: Image.Image):
# 1) detect bboxes
bbox_bg = get_face_bbox(background)
bbox_src = get_face_bbox(source)
if bbox_bg is None:
return None, "❌ No face in background."
if bbox_src is None:
return None, "❌ No face in 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
)
# 3) compute offset
offset = compute_offset(bbox_bg, bbox_src, scale)
# 4) baldhead background
bg_bald = inference(background)
# 5) extract hair-only from source
hair_only = extract_hair(src_scaled)
# 6) paste onto bald background
result = paste_with_alpha(
np.array(bg_bald.convert("RGBA")),
np.array(hair_only),
offset
)
return result
|