Update app.py

app.py

@@ -1,11 +1,11 @@
-import os, json, tempfile, re, traceback
+import os, json, tempfile, re
 import cv2, numpy as np, gradio as gr
 from PIL import Image
 
-# -------- Paths --------
+# -------------------- Paths --------------------
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 CANDIDATES = [
-    os.path.join(BASE_DIR, "hair"),
+    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"),
@@ -17,18 +17,19 @@ if HAIR_DIR is None:
 
 META_PATH = os.path.join(HAIR_DIR, "meta.json")  # optional per-style anchors
 
-# -------- Deps --------
+# -------------------- 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 --------
+# -------------------- Helpers --------------------
 def natural_key(s: str):
-    """Sort 'photo1.png'...'photo11.png' in numeric order."""
+    # 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():
@@ -50,7 +51,7 @@ def load_meta():
     return {}
 
 def premultiply_alpha(bgra):
-    """Removes gray/white halos on edges."""
+    """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)
@@ -59,8 +60,8 @@ def premultiply_alpha(bgra):
 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.ndim != 3 or hair.shape[2] != 4:
-        raise ValueError(f"Invalid hair asset: {name} (must be RGBA PNG in {HAIR_DIR})")
+    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):
@@ -76,14 +77,25 @@ def detect_face_keypoints(img_bgr):
         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):
-    """Return 3 anchor points (L-eye, R-eye, mid-forehead) in hair PNG space."""
     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 perfection, put 3 points per file in hair/meta.json.
+    # 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)
@@ -92,56 +104,66 @@ def hair_reference_points(hair_bgra, filename, meta):
 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
-    # CONSTANT borders avoid clipping/edge artifacts
-    hair_warp = cv2.warpAffine(hair_rgb, M, (W, H), flags=cv2.INTER_LINEAR,
-                               borderMode=cv2.BORDER_CONSTANT, borderValue=(0,0,0))
-    a_warp = cv2.warpAffine(hair_a, M, (W, H), flags=cv2.INTER_LINEAR,
-                            borderMode=cv2.BORDER_CONSTANT, borderValue=0)
+    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, dx, dy, opacity, meta):
-    """Simple engine: NO head-mask (prevents neck lines + lost hair)."""
-    try:
-        if image is None:
-            return None, "Upload a photo first."
-        if not hairstyle:
-            return np.array(image), "Pick a hairstyle."
-
-        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. Use a brighter, front-facing photo."
-
-        hair = load_hair_png(hairstyle)
-        hair_pts = hair_reference_points(hair, hairstyle, meta)
-
-        # destination landmarks + nudges
-        dst = kpts.copy()
-        dst[:, 0] += dx
-        dst[:, 1] += dy
-
-        # scale (no rotation → simpler UI)
-        center = hair_pts.mean(axis=0)
-        s = max(0.5, scale_pct / 100.0)
-        hair_pts_adj = (hair_pts - center) * s + center
-
-        M, _ = cv2.estimateAffinePartial2D(hair_pts_adj, dst, method=cv2.LMEDS)
-        if M is None:
-            return image, "Alignment failed for this image/style."
-
-        out = warp_and_alpha_blend(img_bgr, hair, M, opacity=opacity)
-        return cv2.cvtColor(out, cv2.COLOR_BGR2RGB), "OK"
-    except Exception as e:
-        print("apply_tryon error:", e)
-        traceback.print_exc()
-        return np.array(image) if image is not None else None, "Error applying style."
+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 a hairstyle or 'Apply' first.")
+        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)
@@ -149,13 +171,13 @@ def save_png_to_tmp(img, filename="output_tryon.png"):
         img.save(out_path)
     return out_path
 
-# White thumbnails with captions (number + filename)
+# ---------- 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)
+    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)
@@ -165,127 +187,140 @@ def build_gallery_items(files):
     for idx, fname in enumerate(files, start=1):
         try:
             img = load_hair_png(fname)
-            items.append((thumb_on_white(img), f"{idx}. {fname}"))
+            items.append((thumb_on_white(img), f"{idx}. {fname}"))  # caption shows number & filename
         except Exception:
             continue
     return items
 
-def event_to_filename(evt, files):
-    """Robust across Gradio versions."""
-    idx = getattr(evt, "index", None)
-    if isinstance(idx, int) and 0 <= idx < len(files):
-        return files[idx]
-    val = getattr(evt, "value", None)
-    cap = None
-    if isinstance(val, str):
-        cap = val
-    elif isinstance(val, (tuple, list)) and len(val) >= 2 and isinstance(val[1], str):
-        cap = val[1]
-    elif isinstance(val, dict) and "caption" in val and isinstance(val["caption"], str):
-        cap = val["caption"]
-    if cap:
-        m = re.search(r"\b(\d+)\.\s*(.+)$", cap.strip())
-        if m:
-            i = int(m.group(1)) - 1
-            if 0 <= i < len(files):
-                return files[i]
-        for f in files:
-            if f.lower() == cap.lower():
-                return f
-    return files[0] if files else None
-
-# -------- UI --------
+# -------------------- UI --------------------
 def build_ui():
     META = load_meta()
     HAIR_FILES = load_hairstyles()
 
-    with gr.Blocks(title="Salon Hairstyle Virtual Try-On — Simple & Stable") as demo:
-        gr.Markdown("Upload a photo, then **click a hairstyle**. Adjust size/position if needed. Use **Save result** to export.")
+    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**.")
 
-        selected_file = gr.State(None)
-        meta_state    = gr.State(META)
-        files_state   = gr.State(HAIR_FILES)
+        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():
-                    apply_btn = gr.Button("✨ Apply (optional)")
-                    save_btn  = gr.Button("💾 Save result")
+                    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 styles")
-
+                    refresh = gr.Button("🔄 Refresh")
                 count_md = gr.Markdown(f"Found {len(HAIR_FILES)} hairstyles.")
                 gallery = gr.Gallery(
-                    label="Hairstyles (click to apply)",
+                    label="Hairstyles (click to choose)",
                     value=build_gallery_items(HAIR_FILES),
-                    columns=6, rows=3, height=520,
+                    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 (simple)", open=True):
+                with gr.Accordion("Fine-tune placement", open=True):
                     with gr.Row():
-                        scale = gr.Slider(50, 200, 100, 1, label="Scale (temple distance %)")
-                        opacity = gr.Slider(0.4, 1.0, 1.0, 0.05, label="Hair opacity")
+                        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 (px)")
-                        dy    = gr.Slider(-200, 200, 0, 1, label="Up ↕ Down (px)")
-
-                status = gr.Markdown("")
-
-                def do_apply(im, hairfile, s, dxv, dyv, op, meta):
-                    return apply_tryon(im, hairfile, s, dxv, dyv, op, meta)
-
-                # Click a tile → apply (robust to Gradio variants)
-                def on_gallery_select(evt, files, im, s, dxv, dyv, op, meta):
-                    try:
-                        hairfile = event_to_filename(evt, files)
-                        if hairfile is None:
-                            return None, im, "No styles found."
-                        out, msg = do_apply(im, hairfile, s, dxv, dyv, op, meta)
-                        return hairfile, out, msg
-                    except Exception as e:
-                        print("gallery.select error:", e)
-                        traceback.print_exc()
-                        return None, im, "Error applying style."
-
-                gallery.select(
-                    on_gallery_select,
-                    inputs=[files_state, in_img, scale, dx, dy, opacity, meta_state],
-                    outputs=[selected_file, out_img, status]
-                )
+                        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 after slider tweaks
                 apply_btn.click(
                     fn=do_apply,
-                    inputs=[in_img, selected_file, scale, dx, dy, opacity, meta_state],
+                    inputs=[in_img, hair_sel, scale, rot, dx, dy, opacity, meta_state, limit_head, expand],
                     outputs=[out_img, status]
                 )
 
-                # Save to a real file (shows a link you can click to download)
                 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])
 
-                # Refresh styles list
                 def do_refresh():
                     files = load_hairstyles()
                     items = build_gallery_items(files)
                     msg = f"Found {len(files)} hairstyles."
-                    return items, files, msg
+                    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)
 
-                refresh.click(fn=do_refresh, inputs=[], outputs=[gallery, files_state, count_md])
+                save_live_btn.click(fn=save_snap, inputs=[state_live], outputs=[save_live_file])
 
         return demo
 
-# Spaces autostart
+# Export for Spaces autostart
 app = build_ui()
 demo = app