Upload 41 files

.gitattributes

@@ -33,3 +33,36 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+example_wigs/Heart/HH02.png filter=lfs diff=lfs merge=lfs -text
+example_wigs/Heart/HH03.png filter=lfs diff=lfs merge=lfs -text
+example_wigs/Heart/Loire.png filter=lfs diff=lfs merge=lfs -text
+example_wigs/Heart/SantaRossa.png filter=lfs diff=lfs merge=lfs -text
+example_wigs/Heart/Tuscany.png filter=lfs diff=lfs merge=lfs -text
+example_wigs/Oblong/HH01.png filter=lfs diff=lfs merge=lfs -text
+example_wigs/Oblong/HH02.png filter=lfs diff=lfs merge=lfs -text
+example_wigs/Oblong/HH03.png filter=lfs diff=lfs merge=lfs -text
+example_wigs/Oblong/HH07.png filter=lfs diff=lfs merge=lfs -text
+example_wigs/Oblong/Loire.png filter=lfs diff=lfs merge=lfs -text
+example_wigs/Oval/Alsace.png filter=lfs diff=lfs merge=lfs -text
+example_wigs/Oval/Barossa.png filter=lfs diff=lfs merge=lfs -text
+example_wigs/Oval/Burgundy.png filter=lfs diff=lfs merge=lfs -text
+example_wigs/Oval/HH01.png filter=lfs diff=lfs merge=lfs -text
+example_wigs/Oval/HH02.png filter=lfs diff=lfs merge=lfs -text
+example_wigs/Oval/HH03.png filter=lfs diff=lfs merge=lfs -text
+example_wigs/Oval/HH07.png filter=lfs diff=lfs merge=lfs -text
+example_wigs/Oval/Loire.png filter=lfs diff=lfs merge=lfs -text
+example_wigs/Oval/Napa.png filter=lfs diff=lfs merge=lfs -text
+example_wigs/Oval/Piemonte.png filter=lfs diff=lfs merge=lfs -text
+example_wigs/Oval/Rhone.png filter=lfs diff=lfs merge=lfs -text
+example_wigs/Oval/SantaRossa.png filter=lfs diff=lfs merge=lfs -text
+example_wigs/Oval/Sonoma.png filter=lfs diff=lfs merge=lfs -text
+example_wigs/Oval/Tuscany.png filter=lfs diff=lfs merge=lfs -text
+example_wigs/Round/Loire.png filter=lfs diff=lfs merge=lfs -text
+example_wigs/Round/Piemonte.png filter=lfs diff=lfs merge=lfs -text
+example_wigs/Round/Sonoma.png filter=lfs diff=lfs merge=lfs -text
+example_wigs/Round/Tuscany.png filter=lfs diff=lfs merge=lfs -text
+example_wigs/Square/HH03.png filter=lfs diff=lfs merge=lfs -text
+example_wigs/Square/Loire.png filter=lfs diff=lfs merge=lfs -text
+example_wigs/Square/Piemonte.png filter=lfs diff=lfs merge=lfs -text
+example_wigs/Square/Sonoma.png filter=lfs diff=lfs merge=lfs -text
+example_wigs/Square/Tuscany.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,12 +1,12 @@
----
-title: Be Rejection
-emoji: 🌍
-colorFrom: yellow
-colorTo: purple
-sdk: gradio
-sdk_version: 5.33.2
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+---
+title: Ghep Image
+emoji: 📉
+colorFrom: pink
+colorTo: blue
+sdk: gradio
+sdk_version: 5.31.0
+app_file: app.py
+pinned: false
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,158 @@
+
+import gradio as gr
+from overlay import overlay_source
+from detect_face import predict, NUM_CLASSES
+import os
+from pathlib import Path
+
+BASE_DIR = Path(__file__).parent  # thư mục chứa app.py
+FOLDER = BASE_DIR / "example_wigs"
+
+# --- Hàm load ảnh từ folder ---
+def load_images_from_folder(folder_path: str) -> list[str]:
+    """
+    Trả về list[str] chứa tất cả các hình (jpg, png, gif, bmp) trong folder_path.
+    """
+    supported = {'.jpg', '.jpeg', '.png', '.gif', '.bmp'}
+    if not os.path.isdir(folder_path):
+        print(f"Cảnh báo: '{folder_path}' không phải folder hợp lệ.")
+        return []
+    files = [
+        os.path.join(folder_path, fn)
+        for fn in os.listdir(folder_path)
+        if os.path.splitext(fn)[1].lower() in supported
+    ]
+    if not files:
+        print(f"Không tìm thấy hình trong: {folder_path}")
+    return files
+
+# --- Handler khi click thumbnail của Gallery ---
+# def on_gallery_select(evt: gr.SelectData):
+#     """
+#     Xử lý khi click vào ảnh trong gallery - tối ưu và robust.
+#     """
+#     val = evt.value
+#     if isinstance(val, dict):
+#         img = val.get("image")
+#         if isinstance(img, str): return img
+#         if isinstance(img, dict):
+#             path = img.get("path") or img.get("url")
+#             if isinstance(path, str): return path
+#             for v in img.values():
+#                 if isinstance(v, str) and os.path.isfile(v):
+#                     return v
+#         for v in val.values():
+#             if isinstance(v, str) and os.path.isfile(v):
+#                 return v
+#         raise ValueError(f"Không trích được filepath từ dict: {val}")
+#     if isinstance(val, str):
+#         return val
+#     raise ValueError(f"Kiểu không hỗ trợ: {type(val)}")
+
+def on_gallery_select(evt: gr.SelectData):
+    """
+    Khi click thumbnail: trả về
+      1) filepath để nạp vào Image Source
+      2) tên file (basename) để hiển thị trong Textbox
+    """
+    val = evt.value
+
+    # --- logic trích filepath y như cũ ---
+    if isinstance(val, dict):
+        img = val.get("image")
+        if isinstance(img, str):
+            filepath = img
+        elif isinstance(img, dict):
+            filepath = img.get("path") or img.get("url")
+        else:
+            filepath = next(
+                (v for v in val.values() if isinstance(v, str) and os.path.isfile(v)),
+                None
+            )
+    elif isinstance(val, str):
+        filepath = val
+    else:
+        raise ValueError(f"Kiểu không hỗ trợ: {type(val)}")
+
+    filename = os.path.basename(filepath) if filepath else ""
+    return filepath, filename
+
+# --- Hàm xác định folder dựa trên phân lớp ---
+def infer_folder(image) -> str:
+    cls = predict(image)["predicted_class"]
+    folder = str(FOLDER / cls)
+    return folder
+
+# --- Hàm gộp: phân loại + load ảnh ---
+def handle_bg_change(image):
+    """
+    Khi thay đổi background:
+    1. Phân loại khuôn mặt
+    2. Load ảnh từ folder tương ứng
+    """
+    if image is None:
+        return "", []
+    
+    try:
+        folder = infer_folder(image)
+        images = load_images_from_folder(folder)
+        return folder, images
+    except Exception as e:
+        print(f"Lỗi xử lý ảnh: {e}")
+        return "", []
+
+# --- Xây dựng giao diện Gradio ---
+def build_demo():
+    with gr.Blocks(title="Xử lý hai hình ảnh", theme=gr.themes.Soft()) as demo:
+        gr.Markdown("Upload Background & Source, click **Run** to try on wigs.")
+
+        with gr.Row():
+            bg = gr.Image(type="pil", label="Background", height=500)
+            src = gr.Image(type="pil", label="Source", height=500, interactive=False)
+            out = gr.Image(label="Result", height=500, interactive=False)
+
+        folder_path_box = gr.Textbox(label="Folder path", visible=False)
+
+        
+        with gr.Row():
+            src_name_box = gr.Textbox(
+                        label="Wigs Name",
+                        interactive=False,
+                        show_copy_button=True ,     # tuỳ chọn – tiện copy đường dẫn
+                        scale = 1
+                        )
+            gallery = gr.Gallery(
+                            label="Recommend For You",
+                            height=300,
+                            value=[],
+                            type="filepath",
+                            interactive=False,
+                            columns=5,
+                            object_fit="cover",
+                            allow_preview=True,
+                            scale = 8
+                            )
+            btn = gr.Button("🔄 Run", variant="primary",scale = 1)  
+ 
+
+
+        # Chạy ghép tóc
+        btn.click(fn=overlay_source, inputs=[bg, src], outputs=[out])
+        # Khi đổi ảnh background, tự động phân loại và load ảnh gợi ý
+        bg.change(
+            fn=handle_bg_change,
+            inputs=[bg],
+            outputs=[folder_path_box, gallery],
+            show_progress=True
+        )
+        # Nút tải lại ảnh thủ công (backup)
+        # Khi chọn ảnh trong gallery, cập nhật vào khung Source
+        gallery.select(
+            fn=on_gallery_select,
+            outputs=[src, src_name_box]
+        )
+
+    return demo
+
+if __name__ == "__main__":
+    build_demo().launch()

baldhead.py

@@ -0,0 +1,272 @@
+# baldhead.py
+
+import os
+import cv2
+import numpy as np
+from PIL import Image
+import tensorflow as tf
+import gradio as gr
+
+# Keras imports (note: keras-contrib must be installed)
+import keras.backend as K
+from keras.layers import (
+    Input,
+    Conv2D,
+    UpSampling2D,
+    LeakyReLU,
+    GlobalAveragePooling2D,
+    Dense,
+    Reshape,
+    Dropout,
+    Concatenate,
+    multiply,                 # ← Thêm import multiply
+)
+from keras.models import Model
+from keras_contrib.layers.normalization.instancenormalization import InstanceNormalization
+
+# RetinaFace + skimage for face alignment
+from retinaface import RetinaFace
+from skimage import transform as trans
+
+# Hugging Face Hub helper
+from huggingface_hub import hf_hub_download
+
+
+
+# --- Face‐alignment helpers (giống code gốc) ---
+image_size = [256, 256]
+src_landmarks = np.array([
+    [30.2946, 51.6963],
+    [65.5318, 51.5014],
+    [48.0252, 71.7366],
+    [33.5493, 92.3655],
+    [62.7299, 92.2041],
+], dtype=np.float32)
+src_landmarks[:, 0] += 8.0
+src_landmarks[:, 0] += 15.0
+src_landmarks[:, 1] += 30.0
+src_landmarks /= 112
+src_landmarks *= 200
+
+
+def list2array(values):
+    return np.array(list(values))
+
+
+def align_face(img: np.ndarray):
+    """
+    Detect faces + landmarks in `img` via RetinaFace.
+    Returns lists of aligned face patches (256×256 RGB),
+    corresponding binary masks, and the transformation matrices.
+    """
+    faces = RetinaFace.detect_faces(img)
+    bboxes = np.array([list2array(faces[f]['facial_area']) for f in faces])
+    landmarks = np.array([list2array(faces[f]['landmarks'].values()) for f in faces])
+
+    white_canvas = np.ones(img.shape, dtype=np.uint8) * 255
+    aligned_faces, masks, matrices = [], [], []
+
+    if bboxes.shape[0] > 0:
+        for i in range(bboxes.shape[0]):
+            dst = landmarks[i]  # detected landmarks
+            tform = trans.SimilarityTransform()
+            tform.estimate(dst, src_landmarks)
+            M = tform.params[0:2, :]
+
+            warped_face = cv2.warpAffine(
+                img, M, (image_size[1], image_size[0]), borderValue=0.0
+            )
+            warped_mask = cv2.warpAffine(
+                white_canvas, M, (image_size[1], image_size[0]), borderValue=0.0
+            )
+
+            aligned_faces.append(warped_face)
+            masks.append(warped_mask)
+            matrices.append(tform.params[0:3, :])
+
+    return aligned_faces, masks, matrices
+
+
+def put_face_back(
+    orig_img: np.ndarray,
+    processed_faces: list[np.ndarray],
+    masks: list[np.ndarray],
+    matrices: list[np.ndarray],
+):
+    """
+    Warp each processed face back onto the original `orig_img`
+    using the inverse of the transformation matrices.
+    """
+    result = orig_img.copy()
+    h, w = orig_img.shape[:2]
+
+    for i in range(len(processed_faces)):
+        invM = np.linalg.inv(matrices[i])[0:2]
+        warped = cv2.warpAffine(processed_faces[i], invM, (w, h), borderValue=0.0)
+        mask = cv2.warpAffine(masks[i], invM, (w, h), borderValue=0.0)
+        binary_mask = (mask // 255).astype(np.uint8)
+
+        # Composite: result = result * (1 - mask) + warped * mask
+        result = result * (1 - binary_mask)
+        result = result.astype(np.uint8)
+        result = result + warped * binary_mask
+
+    return result
+
+
+# ----------------------------
+# 2. GENERATOR ARCHITECTURE
+# ----------------------------
+
+def squeeze_excite_block(x, ratio=4):
+    """
+    Squeeze-and-Excitation block: channel-wise attention.
+    """
+    init = x
+    channel_axis = 1 if K.image_data_format() == "channels_first" else -1
+    filters = init.shape[channel_axis]
+    se_shape = (1, 1, filters)
+
+    se = GlobalAveragePooling2D()(init)
+    se = Reshape(se_shape)(se)
+    se = Dense(filters // ratio, activation="relu", kernel_initializer="he_normal", use_bias=False)(se)
+    se = Dense(filters, activation="sigmoid", kernel_initializer="he_normal", use_bias=False)(se)
+    return multiply([init, se])
+
+
+def conv2d(layer_input, filters, f_size=4, bn=True, se=False):
+    """
+    Downsampling block: Conv2D → LeakyReLU → (InstanceNorm) → (SE block)
+    """
+    d = Conv2D(filters, kernel_size=f_size, strides=2, padding="same")(layer_input)
+    d = LeakyReLU(alpha=0.2)(d)
+    if bn:
+        d = InstanceNormalization()(d)
+    if se:
+        d = squeeze_excite_block(d)
+    return d
+
+
+def atrous(layer_input, filters, f_size=4, bn=True):
+    """
+    Atrous (dilated) convolution block with dilation rates [2,4,8].
+    """
+    a_list = []
+    for rate in [2, 4, 8]:
+        a = Conv2D(filters, f_size, dilation_rate=rate, padding="same")(layer_input)
+        a_list.append(a)
+    a = Concatenate()(a_list)
+    a = LeakyReLU(alpha=0.2)(a)
+    if bn:
+        a = InstanceNormalization()(a)
+    return a
+
+
+def deconv2d(layer_input, skip_input, filters, f_size=4, dropout_rate=0):
+    """
+    Upsampling block: UpSampling2D → Conv2D → (Dropout) → InstanceNorm → Concatenate(skip)
+    """
+    u = UpSampling2D(size=2)(layer_input)
+    u = Conv2D(filters, kernel_size=f_size, strides=1, padding="same", activation="relu")(u)
+    if dropout_rate:
+        u = Dropout(dropout_rate)(u)
+    u = InstanceNormalization()(u)
+    u = Concatenate()([u, skip_input])
+    return u
+
+
+def build_generator():
+    """
+    Reconstruct the generator architecture exactly as in the notebook,
+    then return a Keras Model object.
+    """
+    d0 = Input(shape=(256, 256, 3))
+    gf = 64
+
+    # Downsampling
+    d1 = conv2d(d0, gf, bn=False, se=True)
+    d2 = conv2d(d1, gf * 2, se=True)
+    d3 = conv2d(d2, gf * 4, se=True)
+    d4 = conv2d(d3, gf * 8)
+    d5 = conv2d(d4, gf * 8)
+
+    # Atrous block
+    a1 = atrous(d5, gf * 8)
+
+    # Upsampling
+    u3 = deconv2d(a1, d4, gf * 8)
+    u4 = deconv2d(u3, d3, gf * 4)
+    u5 = deconv2d(u4, d2, gf * 2)
+    u6 = deconv2d(u5, d1, gf)
+
+    # Final upsample + conv
+    u7 = UpSampling2D(size=2)(u6)
+    output_img = Conv2D(3, kernel_size=4, strides=1, padding="same", activation="tanh")(u7)
+
+    model = Model(d0, output_img)
+    return model
+
+
+# ----------------------------
+# 3. LOAD MODEL WEIGHTS
+# ----------------------------
+
+HF_REPO_ID = "VanNguyen1214/baldhead"
+HF_FILENAME = "model_G_5_170.hdf5"
+HF_TOKEN = os.environ["HUGGINGFACEHUB_API_TOKEN"]
+
+def load_generator_from_hub():
+    """
+    Download the .hdf5 weights from HF Hub into cache,
+    rebuild the generator, then load weights.
+    """
+    local_path = hf_hub_download(repo_id=HF_REPO_ID, filename=HF_FILENAME,token=HF_TOKEN)
+    gen = build_generator()
+    gen.load_weights(local_path)
+    return gen
+
+# Load once at startup
+try:
+    GENERATOR = load_generator_from_hub()
+    print(f"[INFO] Loaded generator weights from {HF_REPO_ID}/{HF_FILENAME}")
+except Exception as e:
+    print("[ERROR] Could not load generator:", e)
+    GENERATOR = None
+
+
+# ----------------------------
+# 4. INFERENCE FUNCTION
+# ----------------------------
+
+def inference(image: Image.Image) -> Image.Image:
+    """
+    Gradio-compatible inference function:
+    - Convert PIL→ numpy RGB
+    - Align faces
+    - For each face: normalize to [-1,1], run through generator, denormalize to uint8
+    - Put processed faces back onto original image
+    - Return full-image PIL
+    """
+    if GENERATOR is None:
+        return image
+
+    orig = np.array(image.convert("RGB"))
+
+    faces, masks, mats = align_face(orig)
+    if len(faces) == 0:
+        return image
+
+    processed_faces = []
+    for face in faces:
+        face_input = face.astype(np.float32)
+        face_input = (face_input / 127.5) - 1.0  # scale to [-1,1]
+        face_input = np.expand_dims(face_input, axis=0)  # (1,256,256,3)
+
+        pred = GENERATOR.predict(face_input)[0]  # (256,256,3) in [-1,1]
+        pred = ((pred + 1.0) * 127.5).astype(np.uint8)
+        processed_faces.append(pred)
+
+    output_np = put_face_back(orig, processed_faces, masks, mats)
+    output_pil = Image.fromarray(output_np)
+
+    return output_pil

bbox_utils.py

@@ -0,0 +1,31 @@
+import numpy as np
+from PIL import Image
+
+def get_bbox_from_alpha(rgba: Image.Image):
+    arr = np.array(rgba)
+    alpha = arr[...,3]
+    ys, xs = np.where(alpha>0)
+    if ys.size == 0:
+        return None
+    x1, x2 = xs.min(), xs.max()
+    y1, y2 = ys.min(), ys.max()
+    return x1, y1, x2, y2
+
+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)

detect_face.py

@@ -0,0 +1,93 @@
+# model.py
+
+import os
+import sys
+import torch
+import torch.nn as nn
+import torchvision
+from torchvision import transforms
+from huggingface_hub import hf_hub_download
+from PIL import Image
+import numpy as np
+
+# --- Cấu hình chung ---
+DEVICE       = "cuda" if torch.cuda.is_available() else "cpu"
+HF_REPO      = "VanNguyen1214/detect_faceshape"  # repo của bạn trên HF Hub
+HF_FILENAME  = "best_model.pth"                  # file ở root của repo
+LOCAL_CKPT   = "models/best_model.pth"           # sẽ lưu tại đây
+CLASS_NAMES  = ['Heart', 'Oblong', 'Oval', 'Round', 'Square']
+NUM_CLASSES  = len(CLASS_NAMES)
+
+# --- Transform cho ảnh trước inference ---
+_TRANSFORM = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                         std =[0.229, 0.224, 0.225]),
+])
+
+def _ensure_checkpoint() -> str:
+    """
+    Kiểm tra xem LOCAL_CKPT đã tồn tại chưa.
+    Nếu chưa, tải best_model.pth từ HF_REPO và lưu vào ./models/
+    """
+    if os.path.exists(LOCAL_CKPT):
+        return LOCAL_CKPT
+
+    try:
+        ckpt_path = hf_hub_download(
+            repo_id=HF_REPO,
+            filename=HF_FILENAME,
+            local_dir="models",
+        )
+        return ckpt_path
+    except Exception as e:
+        print(f"❌ Không tải được model từ HF Hub: {e}")
+        sys.exit(1)
+
+def _load_model(ckpt_path: str) -> torch.nn.Module:
+    """
+    Tái tạo kiến trúc EfficientNet-B4, load state_dict, đưa về eval mode.
+    """
+    # 1) Khởi tạo EfficientNet-B4
+    model = torchvision.models.efficientnet_b4(pretrained=False)
+    in_features = model.classifier[1].in_features
+    model.classifier = nn.Sequential(
+        nn.Dropout(p=0.3, inplace=True),
+        nn.Linear(in_features, NUM_CLASSES)
+    )
+
+    # 2) Load trọng số
+    state = torch.load(ckpt_path, map_location=DEVICE)
+    model.load_state_dict(state)
+
+    # 3) Đưa model về chế độ evaluation
+    return model.to(DEVICE).eval()
+
+# === Build model ngay khi import ===
+_CKPT_PATH = _ensure_checkpoint()
+_MODEL     = _load_model(_CKPT_PATH)
+
+def predict(image: Image.Image) -> dict:
+    """
+    Chức năng inference:
+      - image: numpy array H×W×3 RGB
+      - Trả về dict:
+          {
+            "predicted_class": str,
+            "confidence": float,
+            "probabilities": { class_name: prob, ... }
+          }
+    """
+    # Convert về PIL + transform
+    img = image.convert("RGB")
+    x   = _TRANSFORM(img).unsqueeze(0).to(DEVICE)
+
+    # Inference
+    with torch.no_grad():
+        logits = _MODEL(x)
+        probs  = torch.softmax(logits, dim=1).squeeze().cpu().numpy()
+
+    idx = int(probs.argmax())
+    return {"predicted_class": CLASS_NAMES[idx]}
+    

overlay.py

@@ -0,0 +1,89 @@
+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

requirements.txt

@@ -0,0 +1,14 @@
+gradio==4.44.0
+transformers==4.36.0
+torch==2.1.0
+torchvision==0.16.0
+huggingface-hub==0.19.4
+Pillow==9.2.0
+opencv-python-headless==4.8.1.78
+numpy==1.24.3
+mediapipe==0.10.8
+tensorflow==2.11.0
+keras==2.11.0
+scikit-image==0.20.0
+git+https://github.com/keras-team/keras-contrib.git
+retina-face==0.0.13

segmentation.py

@@ -0,0 +1,31 @@
+from transformers import SegformerImageProcessor, AutoModelForSemanticSegmentation
+from PIL import Image
+import torch
+import torch.nn.functional as F
+import numpy as np
+
+# Load SegFormer for hair segmentation
+processor = SegformerImageProcessor.from_pretrained("VanNguyen1214/get_face_and_hair")
+model     = AutoModelForSemanticSegmentation.from_pretrained("VanNguyen1214/get_face_and_hair")
+
+def extract_hair(image: Image.Image) -> Image.Image:
+    """
+    Return an RGBA image where hair pixels have alpha=255 and
+    all other pixels have alpha=0.
+    """
+    rgb = image.convert("RGB")
+    arr = np.array(rgb)
+    h, w = arr.shape[:2]
+
+    # Segment hair
+    inputs = processor(images=rgb, return_tensors="pt")
+    with torch.no_grad():
+        logits = model(**inputs).logits.cpu()
+    up = F.interpolate(logits, size=(h, w), mode="bilinear", align_corners=False)
+    seg = up.argmax(dim=1)[0].numpy()
+    hair_mask = (seg == 2).astype(np.uint8)
+
+    # Build RGBA
+    alpha = (hair_mask * 255).astype(np.uint8)
+    rgba  = np.dstack([arr, alpha])
+    return Image.fromarray(rgba)