Update app.py

app.py

@@ -1,243 +1,256 @@
-import os, json, tempfile
-import cv2, numpy as np, gradio as gr
+import os, glob, math
+import numpy as np
 from PIL import Image
+import gradio as gr
 
-# ---------------------- Paths (hair/ first) ----------------------
-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 = None
-for p in CANDIDATES:
-    if os.path.isdir(p):
-        HAIR_DIR = p
-        break
-if HAIR_DIR is None:  # create canonical path if nothing exists yet
-    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
-
-# ---------------------- MediaPipe ----------------------
+# Your hairstyles live here:
+HAIR_DIR = "hair"
+
+# Optional heavy deps (installed via requirements.txt)
 try:
+    import cv2
     import mediapipe as mp
-except Exception as e:
-    raise RuntimeError(f"Mediapipe import failed. Check requirements.txt pins. Details: {e}")
-
-mp_face_mesh = mp.solutions.face_mesh
-mp_selfie_seg = mp.solutions.selfie_segmentation
-LM = {"left_eye_outer": 33, "right_eye_outer": 263, "mid_forehead": 10}
-
-# ---------------------- Helpers ----------------------
-def load_hairstyles():
-    try:
-        files = [f for f in os.listdir(HAIR_DIR) if f.lower().endswith(".png")]
-    except FileNotFoundError:
-        files = []
-    files.sort()
-    return files
-
-HAIR_FILES = load_hairstyles()
-
-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 {}
-META = load_meta()
-
-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))
+    MP_AVAILABLE = True
+except Exception:
+    MP_AVAILABLE = False
+
+APP_TITLE = "Salon Hairstyle Virtual Try-On — Pro Demo"
+
+# --------------------------
+# Asset loading / utilities
+# --------------------------
+def list_hairstyle_files():
+    files = sorted(glob.glob(os.path.join(HAIR_DIR, "*.png")))
+    labels = [os.path.basename(f).replace(".png","").replace("_"," ").title() for f in files]
+    thumbs = [f for f in files]  # Gradio Gallery accepts path strings
+    return labels, files, thumbs
+
+def ensure_rgba(im: Image.Image) -> Image.Image:
+    return im if im.mode == "RGBA" else im.convert("RGBA")
+
+def overlay_rgba(base: Image.Image, overlay: Image.Image, x: int, y: int) -> Image.Image:
+    base = ensure_rgba(base)
+    overlay = ensure_rgba(overlay)
+    bw, bh = base.size
+    ow, oh = overlay.size
+    if ow <= 0 or oh <= 0:
+        return base
+    x1, y1 = max(0, x), max(0, y)
+    x2, y2 = min(bw, x + ow), min(bh, y + oh)
+    if x1 >= x2 or y1 >= y2:
+        return base
+    crop = overlay.crop((x1 - x, y1 - y, x2 - x, y2 - y))
+    region = base.crop((x1, y1, x2, y2))
+    region = Image.alpha_composite(region, crop)
+    base.paste(region, (x1, y1))
+    return base
+
+# --------------------------
+# Face detection / anchoring
+# --------------------------
+def detect_face_keypoints(np_img):
+    """
+    Returns dict:
+      bbox: (x1,y1,x2,y2)
+      temples: (lx,ly, rx,ry)
+      forehead_y: int
+    """
+    h, w = np_img.shape[:2]
+
+    def fallback():
+        cx, cy = w // 2, int(h * 0.42)
+        bw, bh = int(w * 0.4), int(h * 0.45)
+        x1, y1 = max(0, cx - bw // 2), max(0, cy - bh // 2)
+        x2, y2 = min(w, x1 + bw), min(h, y1 + bh)
+        return {
+            "bbox": (x1, y1, x2, y2),
+            "temples": (x1, (y1 + y2) // 2, x2, (y1 + y2) // 2),
+            "forehead_y": max(0, int(y1 - 0.1 * (y2 - y1))),
+        }
+
+    if not MP_AVAILABLE:
+        return fallback()
+
+    mpfm = mp.solutions.face_mesh
+    with mpfm.FaceMesh(static_image_mode=True, max_num_faces=1, refine_landmarks=False) as fm:
+        rgb = cv2.cvtColor(np_img, cv2.COLOR_BGR2RGB) if np_img.shape[2] == 3 else np_img[..., :3]
+        res = fm.process(rgb)
         if not res.multi_face_landmarks:
-            return None
+            return fallback()
+
         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):
-    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.float32)
-        mask = cv2.GaussianBlur(mask, (35, 35), 0)
-        return mask
-
-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 hair
-
-def hair_reference_points(hair_bgra, filename):
-    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 (tune via meta.json for perfection)
-    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)
+        xs = np.array([p.x for p in lm]) * w
+        ys = np.array([p.y for p in lm]) * h
+        x1, y1, x2, y2 = int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())
+
+        def safe_idx(i):
+            i = int(i)
+            i = max(0, min(len(lm) - 1, i))
+            return int(xs[i]), int(ys[i])
+
+        lx, ly = safe_idx(127)  # approx left temple
+        rx, ry = safe_idx(356)  # approx right temple
+        _, fy = safe_idx(10)    # approx forehead
+        y1 = max(0, int(y1 - 0.12 * (y2 - y1)))  # include hairline
+        return {
+            "bbox": (x1, y1, x2, y2),
+            "temples": (lx, ly, rx, ry),
+            "forehead_y": fy,
+        }
+
+def place_and_render(base_np, hair_path, scale, lr_shift_pct, ud_shift_pct, rotation_deg):
+    if base_np is None:
+        raise gr.Error("Please upload or capture a photo first.")
+    base = Image.fromarray(base_np.astype("uint8")).convert("RGBA")
+    key = detect_face_keypoints(np.array(base.convert("RGB")))
+    x1, y1, x2, y2 = key["bbox"]
+    lx, ly, rx, ry = key["temples"]
+    forehead_y = key["forehead_y"]
+
+    hair = Image.open(hair_path).convert("RGBA")
+
+    # Scale from temple distance
+    temple_dx = rx - lx
+    temple_dy = ry - ly
+    temple_dist = max(1, (temple_dx ** 2 + temple_dy ** 2) ** 0.5)
+    target_w = max(1, int(temple_dist * 2.0 * scale))
+    ratio = target_w / hair.width
+    target_h = max(1, int(hair.height * ratio))
+    hair_resized = hair.resize((target_w, target_h), Image.LANCZOS)
+
+    # Auto-rotation + manual extra rotation
+    auto_deg = math.degrees(math.atan2(temple_dy, temple_dx))
+    rot_total = auto_deg + rotation_deg
+    hair_resized = hair_resized.rotate(rot_total, expand=True)
+
+    # Anchor at temple midpoint, above forehead
+    midx = int((lx + rx) / 2)
+    anchor_x = int(midx - hair_resized.width / 2)
+    anchor_y = int(forehead_y - hair_resized.height * 0.45)
+
+    # Apply shifts (percent of image size)
+    img_w, img_h = base.size
+    anchor_x += int(lr_shift_pct * img_w / 100.0)  # +right, -left
+    anchor_y += int(ud_shift_pct * img_h / 100.0)  # +down,  -up
+
+    out = overlay_rgba(base, hair_resized, anchor_x, anchor_y).convert("RGB")
     return out
 
-def apply_tryon(image, hairstyle, scale_pct, rot_deg, dx, dy, opacity):
-    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)
-
-    dst = kpts.copy()
-    dst[:, 0] += dx
-    dst[:, 1] += dy
-
-    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)
-
-    head = person_mask(img_bgr)
-    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(img):
-    if img is None:
-        return None
-    p = os.path.join(tempfile.gettempdir(), "tryon_result.png")
-    Image.fromarray(img).save(p)
-    return p
-
-def hair_preview(hairstyle):
-    if not hairstyle:
-        return None
-    hair = load_hair_png(hairstyle)
-    h, w = hair.shape[:2]
-    tile = 16
-    bg = np.kron(
-        ((np.indices((h//tile+1, w//tile+1)).sum(axis=0) % 2) * 64 + 192).astype(np.uint8),
-        np.ones((tile, tile), np.uint8)
-    )[:h, :w]
-    bg_rgb = np.dstack([bg, bg, bg])
-    a = (hair[:, :, 3:4] / 255.0)
-    comp = (a * hair[:, :, :3] + (1 - a) * bg_rgb).astype(np.uint8)
-    comp = cv2.cvtColor(comp, cv2.COLOR_BGR2RGB)
-    return comp
-
-# ---------------------- UI ----------------------
-def build_ui():
-    with gr.Blocks(title="Virtual Try-On (FR1–FR8)", css="""
-    .gradio-container {max-width: 980px; margin: auto;}
-    @media (max-width: 768px){ .gradio-container {padding: 8px;} }
-    """) as demo:
-        gr.Markdown("## Salon Hairstyle Virtual Try-On\nUpload or use webcam, pick a style from **Select Hairstyle**, adjust, then download.")
-
-        if not HAIR_FILES:
-            gr.Markdown("⚠️ **No hairstyle PNGs found.** Upload files into **`hair/`** (or `assets/hairstyles/`) and reload this Space.")
-
-        with gr.Tabs():
-            # -------- Photo (FR1,3–7) --------
-            with gr.Tab("Photo"):
-                with gr.Row():
-                    in_img = gr.Image(label="Upload photo (JPEG/PNG)", sources=["upload"], type="pil")
-                    hair = gr.Dropdown(
-                        choices=HAIR_FILES,
-                        value=(HAIR_FILES[0] if HAIR_FILES else None),
-                        label="Select Hairstyle (from 'hair/')",
-                        interactive=True
+# --------------------------
+# Gradio callbacks
+# --------------------------
+def refresh_assets():
+    labels, files, thumbs = list_hairstyle_files()
+    if not files:
+        return (
+            gr.update(value=None, choices=[]),
+            gr.update(value=None),
+            "No PNGs found in the 'hair/' folder.",
+        )
+    dd = gr.update(choices=labels, value=labels[0])
+    gallery = gr.update(value=[[p, l] for p, l in zip(thumbs, labels)])
+    return dd, gallery, f"Found {len(files)} hairstyles."
+
+def refresh_dropdown_only():
+    labels, files, _ = list_hairstyle_files()
+    if not files:
+        return gr.update(value=None, choices=[])
+    return gr.update(choices=labels, value=labels[0])
+
+def pick_from_gallery(gallery_select, current_dropdown):
+    # When user clicks a gallery item, we set the dropdown to match
+    labels, _, _ = list_hairstyle_files()
+    if not labels:
+        raise gr.Error("No hairstyles available.")
+    if gallery_select is None:
+        return current_dropdown
+    idx = gallery_select[0]
+    idx = max(0, min(len(labels) - 1, idx))
+    return labels[idx]
+
+def apply_from_dropdown(image, hairstyle_label, scale, lr, ud, rot):
+    labels, files, _ = list_hairstyle_files()
+    if not labels:
+        raise gr.Error("No hairstyle assets found.")
+    if hairstyle_label not in labels:
+        raise gr.Error("Choose a hairstyle first.")
+    path = files[labels.index(hairstyle_label)]
+    return place_and_render(image, path, scale, lr, ud, rot)
+
+def save_image(img_np):
+    if img_np is None:
+        raise gr.Error("Nothing to save. Generate a preview first.")
+    out_path = "output_tryon.png"
+    Image.fromarray(img_np).save(out_path)
+    return out_path
+
+# --------------------------
+# UI
+# --------------------------
+with gr.Blocks(fill_height=True, theme=gr.themes.Soft()) as demo:
+    gr.Markdown(f"# {APP_TITLE}")
+    gr.Markdown("Upload a photo or use webcam. Put **transparent PNGs** in the hair/ folder, then click **Refresh**.")
+
+    with gr.Tabs():
+        # -------- Upload tab --------
+        with gr.Tab("📷 Photo (Upload)"):
+            with gr.Row():
+                with gr.Column():
+                    img = gr.Image(label="Photo", sources=["upload"], type="numpy", height=420)
+
+                    hair_dd = gr.Dropdown(label="Hairstyle", choices=[], interactive=True)
+                    refresh = gr.Button("🔄 Refresh")
+                    status = gr.Markdown("Drop PNGs into hair/ and press Refresh.")
+                    gallery = gr.Gallery(
+                        label="Hairstyles (click to choose)",
+                        columns=4,
+                        height=220,
+                        allow_preview=False,
+                        interactive=True,
                     )
-                with gr.Row():
-                    preview = gr.Image(label="Hairstyle Preview", height=260)
-                hair.change(fn=hair_preview, inputs=[hair], outputs=[preview])
-
-                with gr.Accordion("Alignment Controls", open=True):
-                    with gr.Row():
-                        scale = gr.Slider(50, 200, 100, 1, label="Scale %")
-                        rot = gr.Slider(-30, 30, 0, 1, label="Rotate (deg)")
-                    with gr.Row():
-                        dx = gr.Slider(-200, 200, 0, 1, label="Horizontal Nudge (px)")
-                        dy = gr.Slider(-200, 200, 0, 1, label="Vertical Nudge (px)")
-                    opacity = gr.Slider(0.2, 1.0, 1.0, 0.05, label="Hair Opacity")
-
-                out = gr.Image(label="Result Preview")
-                status = gr.Markdown()
-
-                run = gr.Button("Apply (Align & Overlay)")
-                run.click(
-                    fn=lambda im, h, s, r, dxv, dyv, op: apply_tryon(im, h, s, r, dxv, dyv, op),
-                    inputs=[in_img, hair, scale, rot, dx, dy, opacity],
-                    outputs=[out, status]
-                )
-
-                dl = gr.DownloadButton(label="Download Result", file_name="tryon.png")
-                dl.click(fn=save_png, inputs=[out], outputs=[dl])
-
-                gr.Markdown("Share this Space link after you make it public (FR-7).")
-
-            # -------- Webcam (FR2–FR6) --------
-            with gr.Tab("Webcam"):
-                cam = gr.Image(sources=["webcam"], streaming=True, type="pil", label="Enable camera to start")
-                hair2 = gr.Dropdown(choices=HAIR_FILES, value=(HAIR_FILES[0] if HAIR_FILES else None), label="Select Hairstyle")
-                scale2 = gr.Slider(50, 200, 100, 1, label="Scale %")
-                rot2   = gr.Slider(-25, 25, 0, 1, label="Rotate (deg)")
-                dx2    = gr.Slider(-150, 150, 0, 1, label="Horizontal Nudge (px)")
-                dy2    = gr.Slider(-150, 150, 0, 1, label="Vertical Nudge (px)")
-                opacity2 = gr.Slider(0.2, 1.0, 0.95, 0.05, label="Hair Opacity")
-                out2 = gr.Image(label="Live Preview")
-
-                def live(im, h, s, r, dxv, dyv, op):
-                    res, _ = apply_tryon(im, h, s, r, dxv, dyv, op)
-                    return res
-
-                cam.stream(live, inputs=[cam, hair2, scale2, rot2, dx2, dy2, opacity2], outputs=[out2])
-
-        return demo
-
-# Export for Spaces autostart
-app = build_ui()
-demo = app
-
-# Local dev
+
+                    with gr.Accordion("Fine-tune placement", open=False):
+                        scale = gr.Slider(0.6, 2.2, value=1.3, step=0.01, label="Scale (× temple distance)")
+                        lr = gr.Slider(-30, 30, value=0, step=0.5, label="Left ↔ Right shift (%)")
+                        ud = gr.Slider(-30, 30, value=0, step=0.5, label="Up ↕ Down shift (%)")
+                        rot = gr.Slider(-30, 30, value=0, step=0.5, label="Extra rotation (°)")
+
+                    run_dd = gr.Button("✨ Apply")
+
+                with gr.Column():
+                    out = gr.Image(label="Preview", height=480)
+                    save_btn = gr.Button("💾 Save result")
+                    file_out = gr.File(label="Download")
+
+            # Wiring
+            refresh.click(fn=refresh_assets, inputs=None, outputs=[hair_dd, gallery, status])
+            demo.load(fn=refresh_assets, inputs=None, outputs=[hair_dd, gallery, status])
+
+            # Clicking gallery picks the dropdown value; one Apply button only
+            gallery.select(fn=pick_from_gallery, inputs=[gallery, hair_dd], outputs=hair_dd)
+            run_dd.click(fn=apply_from_dropdown, inputs=[img, hair_dd, scale, lr, ud, rot], outputs=out)
+            save_btn.click(fn=save_image, inputs=out, outputs=file_out)
+
+        # -------- Webcam tab --------
+        with gr.Tab("🎥 Webcam (Live Beta)"):
+            gr.Markdown("Live mode processes frames continuously. Keep resolution small on CPU Spaces.")
+            cam = gr.Image(sources=["webcam"], streaming=True, type="numpy", label="Webcam")
+            hair_dd2 = gr.Dropdown(label="Hairstyle", choices=[], interactive=True)
+            scale2 = gr.Slider(0.6, 2.2, value=1.25, step=0.01, label="Scale")
+            lr2 = gr.Slider(-30, 30, value=0, step=0.5, label="Left ↔ Right shift (%)")
+            ud2 = gr.Slider(-30, 30, value=0, step=0.5, label="Up ↕ Down shift (%)")
+            rot2 = gr.Slider(-30, 30, value=0, step=0.5, label="Rotation (°)")
+            out_live = gr.Image(label="Live Preview", interactive=False, height=420)
+
+            def live_process(frame, label, s, lrp, udp, r):
+                labels, files, _ = list_hairstyle_files()
+                if frame is None or not labels or label not in labels:
+                    return frame
+                path = files[labels.index(label)]
+                return place_and_render(frame, path, s, lrp, udp, r)
+
+            cam.stream(fn=live_process, inputs=[cam, hair_dd2, scale2, lr2, ud2, rot2], outputs=out_live)
+            demo.load(fn=refresh_dropdown_only, inputs=None, outputs=hair_dd2)
+
 if __name__ == "__main__":
-    app.launch()
+    demo.launch()