Update app.py

app.py

@@ -2,7 +2,7 @@ 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"),
@@ -17,7 +17,7 @@ if HAIR_DIR is None:
 
 META_PATH = os.path.join(HAIR_DIR, "meta.json")  # optional per-style anchors
 
-# =============== Dependencies ===============
+# ================== Dependencies ==================
 try:
     import mediapipe as mp
 except Exception as e:
@@ -26,8 +26,9 @@ except Exception as e:
 mp_face_mesh = mp.solutions.face_mesh
 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."""
     return [int(t) if t.isdigit() else t.lower() for t in re.split(r"(\d+)", s)]
 
 def load_hairstyles():
@@ -76,12 +77,13 @@ def detect_face_keypoints(img_bgr):
         return np.stack([xy(LM["left_eye_outer"]), xy(LM["right_eye_outer"]), xy(LM["mid_forehead"])])
 
 def hair_reference_points(hair_bgra, filename, meta):
+    """Return 3 reference 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 pixel-perfect fit, add 3 points to meta.json.
+    # Generic default; add per-style anchors in hair/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)
@@ -91,7 +93,7 @@ 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
-    # borderMode CONSTANT avoids odd edge artifacts; value black (transparent)
+    # CONSTANT border avoids edge artifacts; value=(0,0,0) is transparent with alpha=0
     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,
@@ -101,7 +103,7 @@ def warp_and_alpha_blend(base_bgr, hair_bgra, M, opacity=1.0):
     return out
 
 def apply_tryon(image, hairstyle, scale_pct, dx, dy, opacity, meta):
-    """No head-mask (prevents neck lines & cropping)."""
+    """Simplified engine: no head mask (prevents neck lines/cropping)."""
     if image is None:
         return None, "Upload a photo first."
     if not hairstyle:
@@ -120,7 +122,7 @@ def apply_tryon(image, hairstyle, scale_pct, dx, dy, opacity, meta):
     dst[:, 0] += dx
     dst[:, 1] += dy
 
-    # Scale hair anchors around their centroid (no rotation for simplicity)
+    # Scale hair anchors around their centroid (no rotation by default)
     center = hair_pts.mean(axis=0)
     s = max(0.5, scale_pct / 100.0)
     hair_pts_adj = (hair_pts - center) * s + center
@@ -143,7 +145,7 @@ def save_png_to_tmp(img, filename="output_tryon.png"):
         img.save(out_path)
     return out_path
 
-# ---- white thumbnails with labels ----
+# ---------- white thumbnails with captions ----------
 def thumb_on_white(hair_bgra, max_h=220):
     h, w = hair_bgra.shape[:2]
     scale = min(1.0, max_h / h)
@@ -159,20 +161,48 @@ 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}"))  # show number + filename
+            items.append((thumb_on_white(img), f"{idx}. {fname}"))
         except Exception:
             continue
     return items
 
-# =============== UI ===============
+# Robustly extract a filename from a Gallery select event
+def _event_to_filename(evt, files):
+    """
+    Works across Gradio versions:
+    - Prefer evt.index (int)
+    - Else parse evt.value (like '3. photo3.png')
+    - Else fall back to first file
+    """
+    # 1) index path
+    idx = getattr(evt, "index", None)
+    if isinstance(idx, int) and 0 <= idx < len(files):
+        return files[idx]
+    # 2) value path
+    val = getattr(evt, "value", None)
+    if isinstance(val, (str, tuple, list)):
+        cap = val if isinstance(val, str) else (val[1] if len(val) > 1 else "")
+        if isinstance(cap, str):
+            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]
+            # try raw filename
+            if cap in files:
+                return cap
+    # 3) fallback
+    return files[0] if files else None
+
+# ================== UI ==================
 def build_ui():
     META = load_meta()
     HAIR_FILES = load_hairstyles()
 
     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**.")
+        gr.Markdown("Upload a photo, then **click a hairstyle**. Adjust a little if needed, then **Save result**.")
 
-        selected_file = gr.State(None)  # currently selected hairstyle filename
+        selected_file = gr.State(None)      # current hairstyle filename
         meta_state    = gr.State(META)
         files_state   = gr.State(HAIR_FILES)
 
@@ -200,7 +230,7 @@ def build_ui():
 
                 with gr.Accordion("Fine-tune (simple)", open=True):
                     with gr.Row():
-                        scale = gr.Slider(50, 200, 100, 1, label="Scale (temple distance %)")  # main size
+                        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")
                     with gr.Row():
                         dx    = gr.Slider(-200, 200, 0, 1, label="Left ↔ Right (px)")
@@ -212,13 +242,11 @@ def build_ui():
                 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
+                # 1) gallery click → robust filename → 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]
+                    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
 
@@ -228,7 +256,7 @@ def build_ui():
                     outputs=[selected_file, out_img, status]
                 )
 
-                # 2) Apply button (useful after slider tweaks)
+                # 2) Apply (use after slider tweaks)
                 apply_btn.click(
                     fn=do_apply,
                     inputs=[in_img, selected_file, scale, dx, dy, opacity, meta_state],
@@ -247,7 +275,6 @@ def build_ui():
                     files = load_hairstyles()
                     items = build_gallery_items(files)
                     msg = f"Found {len(files)} hairstyles."
-                    # Keep selection if name still exists
                     return items, files, msg
 
                 refresh.click(fn=do_refresh, inputs=[], outputs=[gallery, files_state, count_md])