Update app.py

app.py

@@ -1,271 +1,59 @@
-import os, glob, math
-import numpy as np
-from PIL import Image
 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**."
+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
     )
 
-    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()
+demo.launch()