Update app.py

app.py

@@ -1,59 +1,271 @@
-import gradio as gr
+import os, glob, math
+import numpy as np
 from PIL import Image
-import os
-
-def load_hairstyles():
-    folder = "hairstyles"
-    if not os.path.exists(folder):
-        return []
-    return [
-        Image.open(os.path.join(folder, f)).convert("RGBA")
-        for f in sorted(os.listdir(folder)) if f.endswith(".png")
-    ]
-
-hairstyles = load_hairstyles()
-
-def apply_hairstyle(user_img, style_index, x_offset, y_offset, scale):
-    if user_img is None or not hairstyles:
-        return None
-
-    user_img = user_img.convert("RGBA")
-    base_w, base_h = user_img.size
-
-    hairstyle = hairstyles[style_index]
-
-    # Resize the hairstyle based on scale
-    new_size = (int(base_w * scale), int(hairstyle.height * (base_w * scale / hairstyle.width)))
-    hairstyle = hairstyle.resize(new_size)
-
-    # Create a blank transparent image to position the hairstyle
-    composite = Image.new("RGBA", user_img.size)
-    paste_x = int((base_w - new_size[0]) / 2 + x_offset)
-    paste_y = int(y_offset)
-    composite.paste(hairstyle, (paste_x, paste_y), hairstyle)
-
-    # Overlay it
-    result = Image.alpha_composite(user_img, composite)
-    return result.convert("RGB")
-
-with gr.Blocks() as demo:
-    gr.Markdown("## 💇 Salon Virtual Hairstyle Try-On (Adjustable)")
-
-    with gr.Row():
-        with gr.Column():
-            image_input = gr.Image(type="pil", label="📷 Upload an Image")
-            style_slider = gr.Slider(0, max(len(hairstyles)-1, 0), step=1, label="🎨 Select Hairstyle")
-            x_offset = gr.Slider(-200, 200, value=0, step=1, label="⬅️➡️ Move Left / Right")
-            y_offset = gr.Slider(-200, 200, value=0, step=1, label="⬆️⬇️ Move Up / Down")
-            scale = gr.Slider(0.3, 2.0, value=1.0, step=0.05, label="📏 Scale Hairstyle")
-            apply_btn = gr.Button("✨ Apply Hairstyle")
-        with gr.Column():
-            result_output = gr.Image(label="🔍 Result Preview")
-
-    apply_btn.click(
-        fn=apply_hairstyle,
-        inputs=[image_input, style_slider, x_offset, y_offset, scale],
-        outputs=result_output
+import gradio as gr
+
+# Optional heavy deps (installed via requirements.txt)
+try:
+    import cv2
+    import mediapipe as mp
+    MP_AVAILABLE = True
+except Exception:
+    MP_AVAILABLE = False
+
+APP_TITLE = "Salon Hairstyle Virtual Try-On — Pro Demo"
+HAIR_DIR = "assets/hairstyles"
+
+# --------------------------
+# Asset loading / utilities
+# --------------------------
+def list_hairstyle_files():
+    files = sorted(glob.glob(os.path.join(HAIR_DIR, "*.png")))
+    labels = [os.path.basename(f).replace(".png","").replace("_"," ").title() for f in files]
+    thumbs = []
+    for f in files:
+        thumbs.append(f)  # Gradio Gallery can take path strings
+    return labels, files, thumbs
+
+def ensure_rgba(im: Image.Image) -> Image.Image:
+    return im if im.mode == "RGBA" else im.convert("RGBA")
+
+def overlay_rgba(base: Image.Image, overlay: Image.Image, x: int, y: int) -> Image.Image:
+    base = ensure_rgba(base)
+    overlay = ensure_rgba(overlay)
+    bw, bh = base.size
+    ow, oh = overlay.size
+    if ow <= 0 or oh <= 0:
+        return base
+    x1, y1 = max(0, x), max(0, y)
+    x2, y2 = min(bw, x + ow), min(bh, y + oh)
+    if x1 >= x2 or y1 >= y2:
+        return base
+    crop = overlay.crop((x1 - x, y1 - y, x2 - x, y2 - y))
+    region = base.crop((x1, y1, x2, y2))
+    region = Image.alpha_composite(region, crop)
+    base.paste(region, (x1, y1))
+    return base
+
+# --------------------------
+# Face detection / anchoring
+# --------------------------
+def detect_face_keypoints(np_img):
+    """
+    Returns dict with keys:
+      bbox: (x1,y1,x2,y2)
+      temples: (lx,ly, rx,ry) approximate
+      forehead_y: int (approx y of forehead landmark)
+    Uses MediaPipe FaceMesh when available; otherwise returns center-based heuristic.
+    """
+    h, w = np_img.shape[:2]
+
+    def fallback():
+        cx, cy = w // 2, int(h * 0.42)
+        bw, bh = int(w * 0.4), int(h * 0.45)
+        x1, y1 = max(0, cx - bw // 2), max(0, cy - bh // 2)
+        x2, y2 = min(w, x1 + bw), min(h, y1 + bh)
+        return {
+            "bbox": (x1, y1, x2, y2),
+            "temples": (x1, (y1 + y2) // 2, x2, (y1 + y2) // 2),
+            "forehead_y": max(0, int(y1 - 0.1 * (y2 - y1))),
+        }
+
+    if not MP_AVAILABLE:
+        return fallback()
+
+    mpfm = mp.solutions.face_mesh
+    with mpfm.FaceMesh(static_image_mode=True, max_num_faces=1, refine_landmarks=False) as fm:
+        rgb = cv2.cvtColor(np_img, cv2.COLOR_BGR2RGB) if np_img.shape[2] == 3 else np_img[..., :3]
+        res = fm.process(rgb)
+        if not res.multi_face_landmarks:
+            return fallback()
+
+        lm = res.multi_face_landmarks[0].landmark
+        xs = np.array([p.x for p in lm]) * w
+        ys = np.array([p.y for p in lm]) * h
+        x1, y1, x2, y2 = int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())
+
+        def safe_idx(i):
+            i = int(i)
+            i = max(0, min(len(lm) - 1, i))
+            return int(xs[i]), int(ys[i])
+
+        lx, ly = safe_idx(127)  # approx left temple
+        rx, ry = safe_idx(356)  # approx right temple
+        _, fy = safe_idx(10)    # approx forehead
+        y1 = max(0, int(y1 - 0.12 * (y2 - y1)))  # expand up a bit
+
+        return {
+            "bbox": (x1, y1, x2, y2),
+            "temples": (lx, ly, rx, ry),
+            "forehead_y": fy,
+        }
+
+def place_and_render(base_np, hair_path, scale, x_shift_pct, y_shift_pct, rotation_deg):
+    if base_np is None:
+        raise gr.Error("Please upload or capture a photo first.")
+    base = Image.fromarray(base_np.astype("uint8")).convert("RGBA")
+    key = detect_face_keypoints(np.array(base.convert("RGB")))
+    x1, y1, x2, y2 = key["bbox"]
+    lx, ly, rx, ry = key["temples"]
+    forehead_y = key["forehead_y"]
+
+    hair = Image.open(hair_path).convert("RGBA")
+
+    # Derive scale from temple distance
+    temple_dx = rx - lx
+    temple_dy = ry - ly
+    temple_dist = max(1, (temple_dx ** 2 + temple_dy ** 2) ** 0.5)
+    target_w = max(1, int(temple_dist * 2.0 * scale))  # widen beyond temples
+    ratio = target_w / hair.width
+    target_h = max(1, int(hair.height * ratio))
+    hair_resized = hair.resize((target_w, target_h), Image.LANCZOS)
+
+    # Auto-rotate with temple slope + extra manual rotation
+    auto_deg = math.degrees(math.atan2(temple_dy, temple_dx))
+    rot_total = auto_deg + rotation_deg
+    hair_resized = hair_resized.rotate(rot_total, expand=True)
+
+    # Anchor: horizontally center at temple midpoint; vertically above forehead
+    midx = int((lx + rx) / 2)
+    anchor_x = int(midx - hair_resized.width / 2)
+    anchor_y = int(forehead_y - hair_resized.height * 0.45)
+
+    # Manual shifts (percent of image size)
+    img_w, img_h = base.size
+    anchor_x += int(x_shift_pct * img_w / 100.0)
+    anchor_y += int(y_shift_pct * img_h / 100.0)
+
+    out = overlay_rgba(base, hair_resized, anchor_x, anchor_y).convert("RGB")
+    return out
+
+# --------------------------
+# Gradio callbacks
+# --------------------------
+def refresh_assets():
+    labels, files, thumbs = list_hairstyle_files()
+    if not files:
+        return (
+            gr.update(value=None, choices=[]),                 # dropdown
+            gr.update(value=None),                             # gallery
+            "No PNGs in assets/hairstyles. Add some and press Refresh.",
+        )
+    dd = gr.update(choices=labels, value=labels[0])
+    gallery = gr.update(value=[[p, l] for p, l in zip(thumbs, labels)])
+    return dd, gallery, f"Found {len(files)} hairstyles."
+
+def refresh_dropdown_only():
+    labels, files, _ = list_hairstyle_files()
+    if not files:
+        return gr.update(value=None, choices=[])
+    return gr.update(choices=labels, value=labels[0])
+
+def pick_from_gallery(gallery_select, current_dropdown):
+    # gallery_select: (index, value)
+    labels, _, _ = list_hairstyle_files()
+    if not labels:
+        raise gr.Error("No hairstyles available.")
+    if gallery_select is None:
+        return current_dropdown
+    idx = gallery_select[0]
+    idx = max(0, min(len(labels) - 1, idx))
+    return labels[idx]
+
+def apply_from_dropdown(image, hairstyle_label, scale, xs, ys, rot):
+    labels, files, _ = list_hairstyle_files()
+    if not labels:
+        raise gr.Error("No hairstyle assets found.")
+    if hairstyle_label not in labels:
+        raise gr.Error("Choose a hairstyle first.")
+    path = files[labels.index(hairstyle_label)]
+    return place_and_render(image, path, scale, xs, ys, rot)
+
+def apply_from_gallery(image, gallery_select, scale, xs, ys, rot):
+    labels, files, _ = list_hairstyle_files()
+    if not labels:
+        raise gr.Error("No hairstyle assets found.")
+    if gallery_select is None:
+        raise gr.Error("Select a hairstyle from the gallery.")
+    idx = gallery_select[0]
+    idx = max(0, min(len(files) - 1, idx))
+    return place_and_render(image, files[idx], scale, xs, ys, rot)
+
+def save_image(img_np):
+    if img_np is None:
+        raise gr.Error("Nothing to save. Generate a preview first.")
+    out_path = "output_tryon.png"
+    Image.fromarray(img_np).save(out_path)
+    return out_path
+
+# --------------------------
+# UI
+# --------------------------
+with gr.Blocks(fill_height=True, theme=gr.themes.Soft()) as demo:
+    gr.Markdown(f"# {APP_TITLE}")
+    gr.Markdown(
+        "Upload a photo or use webcam. Put **transparent PNG** hairstyles in `assets/hairstyles/`, then **Refresh**."
     )
 
-demo.launch()
+    with gr.Tabs():
+        # ------------------ Upload tab ------------------
+        with gr.Tab("📷 Photo (Upload)"):
+            with gr.Row():
+                with gr.Column():
+                    img = gr.Image(label="Photo", sources=["upload"], type="numpy", height=420)
+                    hair_dd = gr.Dropdown(label="Hairstyle (Dropdown)", choices=[], interactive=True)
+                    refresh = gr.Button("🔄 Refresh")
+                    status = gr.Markdown("Add PNGs to `assets/hairstyles/` and press Refresh.")
+                    gallery = gr.Gallery(
+                        label="Hairstyles (click to choose)",
+                        columns=4,
+                        height=220,
+                        allow_preview=False,
+                        interactive=True,
+                    )
+
+                    with gr.Accordion("Fine-tune placement", open=False):
+                        scale = gr.Slider(0.6, 2.2, value=1.3, step=0.01, label="Scale (× temple distance)")
+                        xs = gr.Slider(-30, 30, value=0, step=0.5, label="Horizontal shift (% image width)")
+                        ys = gr.Slider(-30, 30, value=-2, step=0.5, label="Vertical shift (% image height)")
+                        rot = gr.Slider(-30, 30, value=0, step=0.5, label="Extra rotation (°)")
+
+                    run_dd = gr.Button("✨ Apply (Dropdown)")
+                    run_ga = gr.Button("✨ Apply (Gallery selection)")
+
+                with gr.Column():
+                    out = gr.Image(label="Preview", height=480)
+                    save_btn = gr.Button("💾 Save result")
+                    file_out = gr.File(label="Download")
+
+            # Wiring
+            refresh.click(fn=refresh_assets, inputs=None, outputs=[hair_dd, gallery, status])
+            demo.load(fn=refresh_assets, inputs=None, outputs=[hair_dd, gallery, status])
+
+            gallery.select(fn=pick_from_gallery, inputs=[gallery, hair_dd], outputs=hair_dd)
+            run_dd.click(fn=apply_from_dropdown, inputs=[img, hair_dd, scale, xs, ys, rot], outputs=out)
+            run_ga.click(fn=apply_from_gallery, inputs=[img, gallery, scale, xs, ys, rot], outputs=out)
+            save_btn.click(fn=save_image, inputs=out, outputs=file_out)
+
+        # ------------------ Webcam tab ------------------
+        with gr.Tab("🎥 Webcam (Live Beta)"):
+            gr.Markdown("Live mode processes frames continuously. For CPU Spaces, keep webcam resolution modest.")
+            cam = gr.Image(sources=["webcam"], streaming=True, type="numpy", label="Webcam")
+            hair_dd2 = gr.Dropdown(label="Hairstyle", choices=[], interactive=True)
+            scale2 = gr.Slider(0.6, 2.2, value=1.25, step=0.01, label="Scale")
+            xs2 = gr.Slider(-30, 30, value=0, step=0.5, label="X shift (%)")
+            ys2 = gr.Slider(-30, 30, value=-2, step=0.5, label="Y shift (%)")
+            rot2 = gr.Slider(-30, 30, value=0, step=0.5, label="Rotation (°)")
+            out_live = gr.Image(label="Live Preview", interactive=False, height=420)
+
+            def live_process(frame, label, s, x, y, r):
+                labels, files, _ = list_hairstyle_files()
+                if frame is None or not labels or label not in labels:
+                    return frame
+                path = files[labels.index(label)]
+                return place_and_render(frame, path, s, x, y, r)
+
+            cam.stream(fn=live_process, inputs=[cam, hair_dd2, scale2, xs2, ys2, rot2], outputs=out_live)
+            # Populate only the dropdown on load in this tab
+            demo.load(fn=refresh_dropdown_only, inputs=None, outputs=hair_dd2)
+
+if __name__ == "__main__":
+    demo.launch()