Update app.py

app.py

@@ -108,4 +108,199 @@ def apply_tryon(image, hairstyle, scale_pct, rot_deg, dx, dy, opacity, meta):
 
     img_bgr = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
 
-    kpt
+    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. We save to /tmp and return that path."""
+    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 filenames
+        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():
+                    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")  # NOTE: no file_name arg
+                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]
+                )
+
+                # Return a *file path* so the browser downloads with that name
+                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", height=360)
+                state_live = gr.State(None)
+                snap = gr.Button("📸 Snapshot")
+                save_live = gr.DownloadButton("⬇️ Download Snapshot")
+
+                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
+
+# Local dev
+if __name__ == "__main__":
+    app.launch()