Spaces:
Sleeping
Sleeping
| import os, json, tempfile, re | |
| import cv2, numpy as np, gradio as gr | |
| from PIL import Image | |
| # -------------------- Paths -------------------- | |
| BASE_DIR = os.path.dirname(os.path.abspath(__file__)) | |
| CANDIDATES = [ | |
| os.path.join(BASE_DIR, "hair"), # your folder | |
| os.path.join(BASE_DIR, "assets", "hairstyles"), | |
| os.path.join(BASE_DIR, "assets", "Hairstyles"), | |
| os.path.join(BASE_DIR, "hairstyles"), | |
| ] | |
| HAIR_DIR = next((p for p in CANDIDATES if os.path.isdir(p)), None) | |
| if HAIR_DIR is None: | |
| HAIR_DIR = os.path.join(BASE_DIR, "hair") | |
| os.makedirs(HAIR_DIR, exist_ok=True) | |
| META_PATH = os.path.join(HAIR_DIR, "meta.json") # optional per-style anchors | |
| # -------------------- Deps -------------------- | |
| try: | |
| import mediapipe as mp | |
| except Exception as e: | |
| raise RuntimeError(f"Mediapipe import failed. Check requirements pins. Details: {e}") | |
| mp_face_mesh = mp.solutions.face_mesh | |
| mp_selfie_seg = mp.solutions.selfie_segmentation # optional (off by default) | |
| LM = {"left_eye_outer": 33, "right_eye_outer": 263, "mid_forehead": 10} | |
| # -------------------- Helpers -------------------- | |
| def natural_key(s: str): | |
| # sorts photo1, photo2, ... photo10 in numeric order | |
| return [int(t) if t.isdigit() else t.lower() for t in re.split(r"(\d+)", s)] | |
| def load_hairstyles(): | |
| try: | |
| files = [f for f in os.listdir(HAIR_DIR) if f.lower().endswith(".png")] | |
| except FileNotFoundError: | |
| files = [] | |
| files.sort(key=natural_key) | |
| return files | |
| def load_meta(): | |
| if os.path.exists(META_PATH): | |
| try: | |
| with open(META_PATH, "r") as f: | |
| m = json.load(f) | |
| return m if isinstance(m, dict) else {} | |
| except Exception: | |
| return {} | |
| return {} | |
| def premultiply_alpha(bgra): | |
| """Reduce gray/white halos on edges for nicer blending.""" | |
| bgr = bgra[:, :, :3].astype(np.float32) / 255.0 | |
| a = (bgra[:, :, 3:4].astype(np.float32) / 255.0) | |
| bgr_pm = (bgr * a * 255.0).astype(np.uint8) | |
| return np.dstack([bgr_pm, bgra[:, :, 3]]) | |
| def load_hair_png(name): | |
| path = os.path.join(HAIR_DIR, name) | |
| hair = cv2.imread(path, cv2.IMREAD_UNCHANGED) # BGRA | |
| if hair is None or hair.shape[2] != 4: | |
| raise ValueError(f"Invalid hair asset: {name} (must be RGBA PNG)") | |
| return premultiply_alpha(hair) | |
| def detect_face_keypoints(img_bgr): | |
| h, w = img_bgr.shape[:2] | |
| with mp_face_mesh.FaceMesh( | |
| static_image_mode=True, max_num_faces=1, refine_landmarks=True, | |
| min_detection_confidence=0.6 | |
| ) as fm: | |
| res = fm.process(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)) | |
| if not res.multi_face_landmarks: | |
| return None | |
| lm = res.multi_face_landmarks[0].landmark | |
| def xy(i): return np.array([lm[i].x*w, lm[i].y*h], dtype=np.float32) | |
| return np.stack([xy(LM["left_eye_outer"]), xy(LM["right_eye_outer"]), xy(LM["mid_forehead"])]) | |
| def person_mask(img_bgr, expand_px=20): | |
| """Optional head mask (OFF by default). We expand+blur to avoid 'neck lines'.""" | |
| with mp_selfie_seg.SelfieSegmentation(model_selection=1) as seg: | |
| rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) | |
| m = seg.process(rgb).segmentation_mask | |
| mask = (m > 0.5).astype(np.uint8) | |
| if expand_px > 0: | |
| k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2*expand_px+1, 2*expand_px+1)) | |
| mask = cv2.dilate(mask, k, iterations=1) | |
| mask = cv2.GaussianBlur(mask.astype(np.float32), (41, 41), 0) | |
| return mask | |
| def hair_reference_points(hair_bgra, filename, meta): | |
| h, w = hair_bgra.shape[:2] | |
| if filename in meta: | |
| pts = np.array(meta[filename], dtype=np.float32) | |
| if pts.shape == (3, 2): | |
| return pts | |
| # Defaults (ok for many styles). For perfect fit, add 3 points per file to meta.json. | |
| pL = np.array([0.30*w, 0.60*h], dtype=np.float32) | |
| pR = np.array([0.70*w, 0.60*h], dtype=np.float32) | |
| pM = np.array([0.50*w, 0.40*h], dtype=np.float32) | |
| return np.stack([pL, pR, pM], axis=0) | |
| def warp_and_alpha_blend(base_bgr, hair_bgra, M, opacity=1.0): | |
| H, W = base_bgr.shape[:2] | |
| hair_rgb = hair_bgra[:, :, :3] | |
| hair_a = hair_bgra[:, :, 3] / 255.0 | |
| hair_warp = cv2.warpAffine(hair_rgb, M, (W, H), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_TRANSPARENT) | |
| a_warp = cv2.warpAffine(hair_a, M, (W, H), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_TRANSPARENT) | |
| a = np.clip(a_warp * opacity, 0, 1)[..., None] | |
| out = (a * hair_warp + (1 - a) * base_bgr).astype(np.uint8) | |
| return out | |
| def apply_tryon(image, hairstyle, scale_pct, rot_deg, dx, dy, opacity, meta, | |
| limit_head=False, expand_pct=3.0): | |
| """ | |
| limit_head=False by default to avoid 'missing hair' and neck lines. | |
| If True, we use an expanded soft head mask. | |
| """ | |
| if image is None: | |
| return None, "Upload a photo or enable webcam." | |
| if not hairstyle: | |
| return np.array(image), "Pick a hairstyle first." | |
| img_bgr = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) | |
| kpts = detect_face_keypoints(img_bgr) | |
| if kpts is None: | |
| return image, "No face detected. Try a brighter, front-facing photo." | |
| hair = load_hair_png(hairstyle) | |
| hair_pts = hair_reference_points(hair, hairstyle, meta) | |
| # Destination points (with user nudges) | |
| dst = kpts.copy() | |
| dst[:, 0] += dx | |
| dst[:, 1] += dy | |
| # Scale + rotate around hair anchor centroid | |
| center = hair_pts.mean(axis=0) | |
| theta = np.deg2rad(rot_deg) | |
| s = max(0.5, scale_pct / 100.0) | |
| R = np.array([[np.cos(theta), -np.sin(theta)], | |
| [np.sin(theta), np.cos(theta)]], dtype=np.float32) | |
| hair_pts_adj = (hair_pts - center) @ R.T * s + center | |
| M, _ = cv2.estimateAffinePartial2D(hair_pts_adj, dst, method=cv2.LMEDS) | |
| if M is None: | |
| return image, "Could not compute alignment for this image/style." | |
| out = warp_and_alpha_blend(img_bgr, hair, M, opacity=opacity) | |
| if limit_head: | |
| H, W = img_bgr.shape[:2] | |
| expand_px = max(8, int(min(H, W) * (expand_pct / 100.0))) # soft expansion | |
| head = person_mask(img_bgr, expand_px=expand_px) # soft & expanded | |
| head3 = head[..., None] | |
| out = (head3 * out + (1 - head3) * img_bgr).astype(np.uint8) | |
| out_rgb = cv2.cvtColor(out, cv2.COLOR_BGR2RGB) | |
| return out_rgb, "OK" | |
| def save_png_to_tmp(img, filename="output_tryon.png"): | |
| """Create a file in /tmp and return the path (used by the Save button).""" | |
| if img is None: | |
| raise gr.Error("No image to save. Click Apply first.") | |
| out_path = os.path.join(tempfile.gettempdir(), filename) | |
| if isinstance(img, np.ndarray): | |
| Image.fromarray(img).save(out_path) | |
| else: | |
| img.save(out_path) | |
| return out_path | |
| # ---------- WHITE background thumbnails (shows filename number) ---------- | |
| def thumb_on_white(hair_bgra, max_h=220): | |
| h, w = hair_bgra.shape[:2] | |
| scale = min(1.0, max_h / h) | |
| hair_bgra = cv2.resize(hair_bgra, (int(w*scale), int(h*scale)), interpolation=cv2.INTER_LINEAR) | |
| h, w = hair_bgra.shape[:2] | |
| bg_rgb = np.full((h, w, 3), 255, dtype=np.uint8) # white background | |
| a = (hair_bgra[:, :, 3:4] / 255.0) | |
| comp = (a * hair_bgra[:, :, :3] + (1 - a) * bg_rgb).astype(np.uint8) | |
| return cv2.cvtColor(comp, cv2.COLOR_BGR2RGB) | |
| def build_gallery_items(files): | |
| items = [] | |
| for idx, fname in enumerate(files, start=1): | |
| try: | |
| img = load_hair_png(fname) | |
| items.append((thumb_on_white(img), f"{idx}. {fname}")) # caption shows number & filename | |
| except Exception: | |
| continue | |
| return items | |
| # -------------------- UI -------------------- | |
| def build_ui(): | |
| META = load_meta() | |
| HAIR_FILES = load_hairstyles() | |
| with gr.Blocks(title="Salon Hairstyle Virtual Try-On", css=""" | |
| .gradio-container {max-width: 1200px; margin:auto;} | |
| @media (max-width: 768px){ .gradio-container {padding: 8px;} } | |
| """) as demo: | |
| gr.Markdown("Upload a photo or use webcam. Put transparent **PNGs** in **`hair/`**, then click **Refresh**.") | |
| files_state = gr.State(HAIR_FILES) # filenames (natural order) | |
| meta_state = gr.State(META) | |
| with gr.Tabs(): | |
| # -------- Photo Tab -------- | |
| with gr.Tab("📷 Photo (Upload)"): | |
| with gr.Row(): | |
| in_img = gr.Image(label="Input photo (JPEG/PNG)", type="pil", height=360, sources=["upload"]) | |
| out_img = gr.Image(label="Preview", height=360) | |
| with gr.Row(): | |
| hair_sel = gr.Dropdown( | |
| choices=HAIR_FILES, | |
| value=(HAIR_FILES[0] if HAIR_FILES else None), | |
| label="Selected hairstyle", | |
| interactive=True | |
| ) | |
| apply_btn = gr.Button("✨ Apply (Align & Overlay)") | |
| # SAVE (replaces Download) | |
| save_btn = gr.Button("💾 Save result") | |
| save_file = gr.File(label="Saved file", visible=False) | |
| status = gr.Markdown() | |
| with gr.Row(): | |
| refresh = gr.Button("🔄 Refresh") | |
| count_md = gr.Markdown(f"Found {len(HAIR_FILES)} hairstyles.") | |
| gallery = gr.Gallery( | |
| label="Hairstyles (click to choose)", | |
| value=build_gallery_items(HAIR_FILES), | |
| columns=6, rows=3, height=520, # up to 18 tiles visible; all 11 will show | |
| allow_preview=False, object_fit="contain", show_label=True | |
| ) | |
| with gr.Accordion("Fine-tune placement", open=True): | |
| with gr.Row(): | |
| scale = gr.Slider(50, 200, 100, 1, label="Scale (≈ temple distance %)") | |
| rot = gr.Slider(-30, 30, 0, 1, label="Extra rotation (°)") | |
| with gr.Row(): | |
| dx = gr.Slider(-200, 200, 0, 1, label="Left ↔ Right shift (px)") | |
| dy = gr.Slider(-200, 200, 0, 1, label="Up ↕ Down shift (px)") | |
| opacity = gr.Slider(0.2, 1.0, 1.0, 0.05, label="Hair opacity") | |
| limit_head = gr.Checkbox(label="Limit overlay to head (avoid spill)", value=False) | |
| expand = gr.Slider(0.0, 10.0, 3.0, 0.5, label="Head-mask expansion (%) — only if enabled") | |
| # --- Callbacks --- | |
| def do_apply(im, hfile, s, r, dxv, dyv, op, meta, lh, ex): | |
| return apply_tryon(im, hfile, s, r, dxv, dyv, op, meta, limit_head=lh, expand_pct=ex) | |
| apply_btn.click( | |
| fn=do_apply, | |
| inputs=[in_img, hair_sel, scale, rot, dx, dy, opacity, meta_state, limit_head, expand], | |
| outputs=[out_img, status] | |
| ) | |
| def do_save(im): | |
| path = save_png_to_tmp(im, "output_tryon.png") | |
| return gr.File.update(value=path, visible=True) | |
| save_btn.click(fn=do_save, inputs=[out_img], outputs=[save_file]) | |
| def do_refresh(): | |
| files = load_hairstyles() | |
| items = build_gallery_items(files) | |
| msg = f"Found {len(files)} hairstyles." | |
| return items, gr.update(choices=files, value=(files[0] if files else None)), files, msg | |
| refresh.click(fn=do_refresh, inputs=[], outputs=[gallery, hair_sel, files_state, count_md]) | |
| # Gallery click -> set dropdown to that filename | |
| def on_gallery_select(evt, files): | |
| idx = getattr(evt, "index", None) | |
| if idx is None or not files: | |
| return gr.update() | |
| # our captions start at 1., map index to filename directly | |
| idx = max(0, min(idx, len(files)-1)) | |
| return gr.update(value=files[idx]) | |
| gallery.select(on_gallery_select, inputs=[files_state], outputs=[hair_sel]) | |
| # -------- Webcam Tab (unchanged except 'Save Snapshot') -------- | |
| with gr.Tab("📹 Webcam (Live Beta)"): | |
| cam = gr.Image(sources=["webcam"], streaming=True, type="pil", label="Enable camera") | |
| hair2 = gr.Dropdown(choices=HAIR_FILES, value=(HAIR_FILES[0] if HAIR_FILES else None), label="Selected hairstyle") | |
| with gr.Row(): | |
| scale2 = gr.Slider(50, 200, 100, 1, label="Scale %") | |
| rot2 = gr.Slider(-25, 25, 0, 1, label="Rotate (°)") | |
| with gr.Row(): | |
| dx2 = gr.Slider(-150, 150, 0, 1, label="Left ↔ Right (px)") | |
| dy2 = gr.Slider(-150, 150, 0, 1, label="Up ↕ Down (px)") | |
| opacity2 = gr.Slider(0.2, 1.0, 0.95, 0.05, label="Hair opacity") | |
| limit_head2 = gr.Checkbox(label="Limit overlay to head", value=False) | |
| expand2 = gr.Slider(0.0, 10.0, 3.0, 0.5, label="Head-mask expansion (%)", visible=True) | |
| out2 = gr.Image(label="Live result", height=360) | |
| state_live = gr.State(None) | |
| snap = gr.Button("📸 Snapshot") | |
| save_live_btn = gr.Button("💾 Save snapshot") | |
| save_live_file = gr.File(label="snapshot", visible=False) | |
| def live(im, h, s, r, dxv, dyv, op, meta, lh, ex): | |
| res, _ = apply_tryon(im, h, s, r, dxv, dyv, op, meta, limit_head=lh, expand_pct=ex) | |
| return res, res | |
| cam.stream( | |
| fn=live, | |
| inputs=[cam, hair2, scale2, rot2, dx2, dy2, opacity2, meta_state, limit_head2, expand2], | |
| outputs=[out2, state_live] | |
| ) | |
| snap.click(lambda x: x, inputs=[state_live], outputs=[out2]) | |
| def save_snap(im): | |
| path = save_png_to_tmp(im, "tryon_webcam.png") | |
| return gr.File.update(value=path, visible=True) | |
| save_live_btn.click(fn=save_snap, inputs=[state_live], outputs=[save_live_file]) | |
| return demo | |
| # Export for Spaces autostart | |
| app = build_ui() | |
| demo = app | |
| if __name__ == "__main__": | |
| app.launch() | |