Spaces:
Runtime error
Runtime error
| import numpy as np | |
| from PIL import Image | |
| import mediapipe as mp | |
| import os | |
| from segmentation import get_face, 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: | |
| """Get face bounding box using MediaPipe""" | |
| 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_to_match_faces(bbox_bg, bbox_src) -> float: | |
| """Compute scale factor to make source face same size as background face""" | |
| 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 to match both width and height | |
| scale_w = w_bg / w_src | |
| scale_h = h_bg / h_src | |
| # Use average scale to maintain aspect ratio | |
| scale = (scale_w + scale_h) / 2 | |
| return scale | |
| def compute_offset_to_align_faces(bbox_bg, bbox_src, scale) -> tuple[int,int]: | |
| """Compute offset to align source face center with background face center""" | |
| # Background face center | |
| x1, y1, x2, y2 = bbox_bg | |
| bg_cx = x1 + (x2 - x1) // 2 | |
| bg_cy = y1 + (y2 - y1) // 2 | |
| # Source face center (after scaling) | |
| sx1, sy1, sx2, sy2 = bbox_src | |
| src_cx = int((sx1 + (sx2 - sx1) // 2) * scale) | |
| src_cy = int((sy1 + (sy2 - sy1) // 2) * scale) | |
| # Offset to align centers | |
| offset_x = bg_cx - src_cx | |
| offset_y = bg_cy - src_cy | |
| return offset_x, offset_y | |
| def paste_with_alpha(bg: np.ndarray, src: np.ndarray, offset: tuple[int,int]) -> Image.Image: | |
| """Paste source image onto background using alpha channel""" | |
| res = bg.copy() | |
| x, y = offset | |
| h, w = src.shape[:2] | |
| # Calculate valid region bounds | |
| 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) | |
| # Get source and destination regions | |
| cs = src[y1-y:y2-y, x1-x:x2-x] | |
| cd = res[y1:y2, x1:x2] | |
| # Apply alpha blending | |
| if cs.shape[2] == 4: # RGBA source | |
| mask = cs[..., 3] > 0 | |
| if cd.shape[2] == 3: # RGB destination | |
| cd[mask] = cs[mask][..., :3] | |
| else: # RGBA destination | |
| cd[mask] = cs[mask] | |
| else: # RGB source | |
| cd[:] = cs | |
| res[y1:y2, x1:x2] = cd | |
| return Image.fromarray(res) | |
| def paste_with_alpha_smooth(bg: np.ndarray, src: np.ndarray, offset: tuple[int,int]) -> Image.Image: | |
| """Paste source image onto background using alpha channel with edge smoothing""" | |
| import cv2 | |
| res = bg.copy() | |
| x, y = offset | |
| h, w = src.shape[:2] | |
| # Calculate valid region bounds | |
| 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) | |
| # Get source and destination regions | |
| cs = src[y1-y:y2-y, x1-x:x2-x] | |
| cd = res[y1:y2, x1:x2] | |
| # Apply alpha blending with smoothing | |
| if cs.shape[2] == 4: # RGBA source | |
| # Get alpha channel and smooth it | |
| alpha = cs[..., 3].astype(np.float32) / 255.0 | |
| # Apply Gaussian blur to alpha for smooth edges | |
| alpha_smooth = cv2.GaussianBlur(alpha, (5, 5), 1.0) | |
| alpha_smooth = np.clip(alpha_smooth, 0, 1) | |
| # Expand alpha to 3 channels for RGB blending | |
| alpha_3ch = np.stack([alpha_smooth] * 3, axis=-1) | |
| if cd.shape[2] == 3: # RGB destination | |
| # Smooth blending | |
| cs_rgb = cs[..., :3].astype(np.float32) | |
| cd_float = cd.astype(np.float32) | |
| blended = cd_float * (1 - alpha_3ch) + cs_rgb * alpha_3ch | |
| cd[:] = np.clip(blended, 0, 255).astype(np.uint8) | |
| else: # RGBA destination | |
| # For RGBA destination, blend RGB and keep alpha | |
| cs_rgb = cs[..., :3].astype(np.float32) | |
| cd_rgb = cd[..., :3].astype(np.float32) | |
| blended_rgb = cd_rgb * (1 - alpha_3ch) + cs_rgb * alpha_3ch | |
| cd[..., :3] = np.clip(blended_rgb, 0, 255).astype(np.uint8) | |
| # Update alpha channel | |
| cd_alpha = cd[..., 3].astype(np.float32) / 255.0 | |
| new_alpha = np.maximum(cd_alpha, alpha_smooth) | |
| cd[..., 3] = (new_alpha * 255).astype(np.uint8) | |
| else: # RGB source | |
| cd[:] = cs | |
| res[y1:y2, x1:x2] = cd | |
| return Image.fromarray(res) | |
| def overlay_source(background: Image.Image, source: Image.Image): | |
| """ | |
| New workflow: | |
| 1. Scale source image so source face = background face size | |
| 2. Get face from background and overlay onto scaled source | |
| 3. Extract hair+face+forehead from result and overlay onto background with smoothing, | |
| then get hair from source and overlay onto temp variable | |
| """ | |
| # Step 1: Detect faces in both images | |
| bbox_bg = get_face_bbox(background) | |
| bbox_src = get_face_bbox(source) | |
| if bbox_bg is None: | |
| print("❌ No face detected in background image") | |
| return None | |
| if bbox_src is None: | |
| print("❌ No face detected in source image") | |
| return None | |
| print(f"✅ Background face bbox: {bbox_bg}") | |
| print(f"✅ Source face bbox: {bbox_src}") | |
| # Step 1: Scale source image so faces are same size | |
| scale = compute_scale_to_match_faces(bbox_bg, bbox_src) | |
| scaled_source = source.resize( | |
| (int(source.width * scale), int(source.height * scale)), | |
| Image.Resampling.LANCZOS | |
| ) | |
| print(f"✅ Scaled source by factor: {scale:.2f}") | |
| # Update source bbox after scaling | |
| bbox_src_scaled = ( | |
| int(bbox_src[0] * scale), | |
| int(bbox_src[1] * scale), | |
| int(bbox_src[2] * scale), | |
| int(bbox_src[3] * scale) | |
| ) | |
| # Step 2: Get face from background (with forehead) and overlay onto scaled source | |
| bg_face = get_face(background) # Extract face with forehead | |
| # Compute offset to align background face onto scaled source face | |
| offset = compute_offset_to_align_faces(bbox_src_scaled, bbox_bg, 1.0) | |
| # Overlay background face onto scaled source | |
| source_with_bg_face = paste_with_alpha( | |
| np.array(scaled_source.convert("RGBA")), | |
| np.array(bg_face), | |
| offset | |
| ) | |
| print("✅ Overlaid background face onto scaled source") | |
| # Step 3: Extract hair+face+forehead from the result | |
| hair_face_result = get_face(source_with_bg_face) | |
| # Overlay this onto original background with smoothing -> temp variable | |
| bbox_result = get_face_bbox(source_with_bg_face) | |
| if bbox_result is None: | |
| print("❌ Could not detect face in intermediate result") | |
| return None | |
| final_offset = compute_offset_to_align_faces(bbox_bg, bbox_result, 1.0) | |
| # Overlay onto background with smoothing -> temp variable | |
| temp_result = paste_with_alpha_smooth( | |
| np.array(background.convert("RGBA")), | |
| np.array(hair_face_result), | |
| final_offset | |
| ) | |
| print("✅ Overlaid hair+face+forehead onto background with smoothing -> temp variable") | |
| # Get hair from source and overlay onto temp variable | |
| source_hair = extract_hair(scaled_source) | |
| # Calculate offset to align source hair with temp_result hair position | |
| # Since source was scaled and positioned, we need to use the same offset as final_offset | |
| hair_offset = final_offset | |
| # Overlay source hair onto temp result with correct position | |
| temp_with_source_hair = paste_with_alpha_smooth( | |
| np.array(temp_result), | |
| np.array(source_hair), | |
| hair_offset # Use same offset to align with temp_result hair position | |
| ) | |
| print("✅ Overlaid source hair onto temp result") | |
| # Return final temp result | |
| return temp_with_source_hair.convert("RGB") | |