Spaces:
Running
Running
| 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 | |