app.py

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()