Update app.py

app.py

@@ -2,10 +2,10 @@ 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"),                      # your folder
+    os.path.join(BASE_DIR, "hair"),
     os.path.join(BASE_DIR, "assets", "hairstyles"),
     os.path.join(BASE_DIR, "assets", "Hairstyles"),
     os.path.join(BASE_DIR, "hairstyles"),
@@ -17,19 +17,17 @@ if HAIR_DIR is None:
 
 META_PATH = os.path.join(HAIR_DIR, "meta.json")  # optional per-style anchors
 
-# -------------------- Deps --------------------
+# =============== Dependencies ===============
 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):
-    # 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():
@@ -51,7 +49,7 @@ def load_meta():
     return {}
 
 def premultiply_alpha(bgra):
-    """Reduce gray/white halos on edges for nicer blending."""
+    """Eliminate gray/white halos on edges."""
     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)
@@ -77,25 +75,13 @@ 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):
     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.
+    # Defaults (OK for many styles). For pixel-perfect fit, add 3 points 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)
@@ -105,65 +91,51 @@ 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)
+    # borderMode CONSTANT avoids odd edge artifacts; value black (transparent)
+    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)
     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.
-    """
+def apply_tryon(image, hairstyle, scale_pct, dx, dy, opacity, meta):
+    """No head-mask (prevents neck lines & cropping)."""
     if image is None:
-        return None, "Upload a photo or enable webcam."
+        return None, "Upload a photo first."
     if not hairstyle:
-        return np.array(image), "Pick a hairstyle first."
+        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. Try a brighter, front-facing photo."
+        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 points (with user nudges)
+    # Target points = facial anchors + user nudges
     dst = kpts.copy()
     dst[:, 0] += dx
     dst[:, 1] += dy
 
-    # Scale + rotate around hair anchor centroid
+    # Scale hair anchors around their centroid (no rotation for simplicity)
     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
+    hair_pts_adj = (hair_pts - center) * 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."
+        return image, "Alignment failed 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.")
+        raise gr.Error("No image to save. Click a hairstyle or 'Apply' first.")
     out_path = os.path.join(tempfile.gettempdir(), filename)
     if isinstance(img, np.ndarray):
         Image.fromarray(img).save(out_path)
@@ -171,13 +143,13 @@ def save_png_to_tmp(img, filename="output_tryon.png"):
         img.save(out_path)
     return out_path
 
-# ---------- WHITE background thumbnails (shows filename number) ----------
+# ---- white thumbnails with labels ----
 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
+    bg_rgb = np.full((h, w, 3), 255, dtype=np.uint8)
     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)
@@ -187,140 +159,102 @@ 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}"))  # caption shows number & filename
+            items.append((thumb_on_white(img), f"{idx}. {fname}"))  # show number + filename
         except Exception:
             continue
     return items
 
-# -------------------- UI --------------------
+# =============== 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**.")
+    with gr.Blocks(title="Salon Hairstyle Virtual Try-On (Simple)") as demo:
+        gr.Markdown("Upload a photo, then **click a hairstyle** below. Use a few sliders if needed, then **Save result**.")
 
-        files_state = gr.State(HAIR_FILES)   # filenames (natural order)
-        meta_state  = gr.State(META)
+        selected_file = gr.State(None)  # currently selected hairstyle filename
+        meta_state    = gr.State(META)
+        files_state   = gr.State(HAIR_FILES)
 
         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")
+                    apply_btn = gr.Button("✨ Apply (optional)")
+                    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")
+                    refresh = gr.Button("🔄 Refresh styles")
+
                 count_md = gr.Markdown(f"Found {len(HAIR_FILES)} hairstyles.")
                 gallery = gr.Gallery(
-                    label="Hairstyles (click to choose)",
+                    label="Hairstyles (click to apply)",
                     value=build_gallery_items(HAIR_FILES),
-                    columns=6, rows=3, height=520,  # up to 18 tiles visible; all 11 will show
+                    columns=6, rows=3, height=520,
                     allow_preview=False, object_fit="contain", show_label=True
                 )
 
-                with gr.Accordion("Fine-tune placement", open=True):
+                with gr.Accordion("Fine-tune (simple)", 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 (°)")
+                        scale = gr.Slider(50, 200, 100, 1, label="Scale (temple distance %)")  # main size
+                        opacity = gr.Slider(0.4, 1.0, 1.0, 0.05, label="Hair opacity")
                     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")
+                        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("")
 
-                # --- 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)
+                # ----- actions -----
+                def do_apply(im, hairfile, s, dxv, dyv, op, meta):
+                    return apply_tryon(im, hairfile, s, dxv, dyv, op, meta)
 
+                # 1) click a tile -> set selected file AND auto-apply
+                def on_gallery_select(evt, files, im, s, dxv, dyv, op, meta):
+                    idx = getattr(evt, "index", None)
+                    if idx is None or not files:
+                        return None, gr.update(), None
+                    idx = max(0, min(idx, len(files)-1))
+                    hairfile = files[idx]
+                    out, msg = do_apply(im, hairfile, s, dxv, dyv, op, meta)
+                    return hairfile, out, msg
+
+                gallery.select(
+                    on_gallery_select,
+                    inputs=[files_state, in_img, scale, dx, dy, opacity, meta_state],
+                    outputs=[selected_file, out_img, status]
+                )
+
+                # 2) Apply button (useful after slider tweaks)
                 apply_btn.click(
                     fn=do_apply,
-                    inputs=[in_img, hair_sel, scale, rot, dx, dy, opacity, meta_state, limit_head, expand],
+                    inputs=[in_img, selected_file, scale, dx, dy, opacity, meta_state],
                     outputs=[out_img, status]
                 )
 
+                # 3) Save
                 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])
 
+                # 4) Refresh styles
                 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)
+                    # Keep selection if name still exists
+                    return items, files, msg
 
-                save_live_btn.click(fn=save_snap, inputs=[state_live], outputs=[save_live_file])
+                refresh.click(fn=do_refresh, inputs=[], outputs=[gallery, files_state, count_md])
 
         return demo
 
-# Export for Spaces autostart
+# export for Spaces
 app = build_ui()
 demo = app