Update app.py

app.py

@@ -4,9 +4,8 @@ from PIL import Image
 
 # ===================== Paths & Assets =====================
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
-# Look for your folder first (hair/). Fall back to other common names.
 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"),
@@ -22,7 +21,7 @@ if HAIR_DIR is None:
 
 META_PATH = os.path.join(HAIR_DIR, "meta.json")  # optional per-style anchors
 
-# ===================== Dependences =====================
+# ===================== Dependencies =====================
 try:
     import mediapipe as mp
 except Exception as e:
@@ -85,7 +84,7 @@ def hair_reference_points(hair_bgra, filename, meta):
         pts = np.array(meta[filename], dtype=np.float32)
         if pts.shape == (3, 2):
             return pts
-    # Defaults (works OK for many styles). For pixel-perfect fit, add meta.json.
+    # Defaults (ok for many styles). For perfect fit, put 3 points per file in 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)
@@ -109,200 +108,4 @@ def apply_tryon(image, hairstyle, scale_pct, rot_deg, dx, dy, opacity, meta):
 
     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)
-
-    # Restrict to head region for cleaner look
-    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_to_tmp(img, filename="output_tryon.png"):
-    """DownloadButton needs a file path string."""
-    if img is None:
-        raise gr.Error("No image to download. 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
-
-# ----- thumbnails on checkerboard for the gallery -----
-def thumb_on_checker(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]
-    tile = 12
-    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_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 fname in files:
-        try:
-            img = load_hair_png(fname)
-            items.append((thumb_on_checker(img), fname))  # (image, caption)
-        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 — Pro Demo", css="""
-    .gradio-container {max-width: 1100px; margin:auto;}
-    @media (max-width: 768px){ .gradio-container {padding: 8px;} }
-    """) as demo:
-        gr.Markdown(
-            "### Upload a photo or use webcam. Put transparent **PNGs** in the **`hair/`** folder, then click **Refresh**."
-        )
-
-        files_state = gr.State(HAIR_FILES)   # keep current filenames order
-        meta_state  = gr.State(META)         # pass meta to callbacks
-
-        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)")
-                    download_btn = gr.DownloadButton("⬇️ Download", file_name="output_tryon.png")
-                status = gr.Markdown()
-
-                with gr.Row():
-                    refresh = gr.Button("🔄 Refresh")
-                gallery = gr.Gallery(
-                    label="Hairstyles (click to choose)",
-                    value=build_gallery_items(HAIR_FILES),
-                    columns=6, rows=2, height=320,
-                    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")
-
-                # --- Callbacks ---
-                def do_apply(im, hfile, s, r, dxv, dyv, op, meta):
-                    return apply_tryon(im, hfile, s, r, dxv, dyv, op, meta)
-
-                apply_btn.click(
-                    fn=do_apply,
-                    inputs=[in_img, hair_sel, scale, rot, dx, dy, opacity, meta_state],
-                    outputs=[out_img, status]
-                )
-
-                download_btn.click(
-                    fn=lambda im: save_png_to_tmp(im, "output_tryon.png"),
-                    inputs=[out_img],
-                    outputs=[download_btn]
-                )
-
-                def do_refresh():
-                    files = load_hairstyles()
-                    items = build_gallery_items(files)
-                    return items, gr.update(choices=files, value=(files[0] if files else None)), files
-
-                refresh.click(fn=do_refresh, inputs=[], outputs=[gallery, hair_sel, files_state])
-
-                # Clicking a tile sets the 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()
-                    if idx >= len(files):
-                        idx = len(files) - 1
-                    return gr.update(value=files[idx])
-
-                gallery.select(on_gallery_select, inputs=[files_state], outputs=[hair_sel])
-
-            # ---------------- Webcam Tab ----------------
-            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")
-                out2 = gr.Image(label="Live result")
-                state_live = gr.State(None)
-                snap = gr.Button("📸 Snapshot")
-                save_live = gr.DownloadButton("⬇️ Download Snapshot", file_name="tryon_webcam.png")
-
-                def live(im, h, s, r, dxv, dyv, op, meta):
-                    res, _ = apply_tryon(im, h, s, r, dxv, dyv, op, meta)
-                    return res, res
-
-                cam.stream(
-                    fn=live,
-                    inputs=[cam, hair2, scale2, rot2, dx2, dy2, opacity2, meta_state],
-                    outputs=[out2, state_live]
-                )
-
-                snap.click(lambda x: x, inputs=[state_live], outputs=[out2])
-
-                save_live.click(
-                    fn=lambda im: save_png_to_tmp(im, "tryon_webcam.png"),
-                    inputs=[state_live],
-                    outputs=[save_live]
-                )
-
-        return demo
-
-# Export for Spaces autostart
-app = build_ui()
-demo = app  # alias
-
-# Local dev
-if __name__ == "__main__":
-    app.launch()
+    kpt