Update app.py

app.py

@@ -1,54 +1,243 @@
-import gradio as gr
+import os, json, tempfile
+import cv2, numpy as np, gradio as gr
 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")
+
+# ---------------------- Paths (hair/ first) ----------------------
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+CANDIDATES = [
+    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"),
 ]
-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
-    )
-
-demo.launch()
+HAIR_DIR = None
+for p in CANDIDATES:
+    if os.path.isdir(p):
+        HAIR_DIR = p
+        break
+if HAIR_DIR is None:  # create canonical path if nothing exists yet
+    HAIR_DIR = os.path.join(BASE_DIR, "hair")
+    os.makedirs(HAIR_DIR, exist_ok=True)
+
+META_PATH = os.path.join(HAIR_DIR, "meta.json")  # optional per-style anchors
+
+# ---------------------- MediaPipe ----------------------
+try:
+    import mediapipe as mp
+except Exception as e:
+    raise RuntimeError(f"Mediapipe import failed. Check requirements.txt pins. Details: {e}")
+
+mp_face_mesh = mp.solutions.face_mesh
+mp_selfie_seg = mp.solutions.selfie_segmentation
+LM = {"left_eye_outer": 33, "right_eye_outer": 263, "mid_forehead": 10}
+
+# ---------------------- Helpers ----------------------
+def load_hairstyles():
+    try:
+        files = [f for f in os.listdir(HAIR_DIR) if f.lower().endswith(".png")]
+    except FileNotFoundError:
+        files = []
+    files.sort()
+    return files
+
+HAIR_FILES = load_hairstyles()
+
+def load_meta():
+    if os.path.exists(META_PATH):
+        try:
+            with open(META_PATH, "r") as f:
+                m = json.load(f)
+            return m if isinstance(m, dict) else {}
+        except Exception:
+            return {}
+    return {}
+META = load_meta()
+
+def detect_face_keypoints(img_bgr):
+    h, w = img_bgr.shape[:2]
+    with mp_face_mesh.FaceMesh(
+        static_image_mode=True, max_num_faces=1, refine_landmarks=True,
+        min_detection_confidence=0.6
+    ) as fm:
+        res = fm.process(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB))
+        if not res.multi_face_landmarks:
+            return None
+        lm = res.multi_face_landmarks[0].landmark
+        def xy(i): return np.array([lm[i].x*w, lm[i].y*h], dtype=np.float32)
+        return np.stack([xy(LM["left_eye_outer"]), xy(LM["right_eye_outer"]), xy(LM["mid_forehead"])])
+
+def person_mask(img_bgr):
+    with mp_selfie_seg.SelfieSegmentation(model_selection=1) as seg:
+        rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
+        m = seg.process(rgb).segmentation_mask
+        mask = (m > 0.5).astype(np.float32)
+        mask = cv2.GaussianBlur(mask, (35, 35), 0)
+        return mask
+
+def load_hair_png(name):
+    path = os.path.join(HAIR_DIR, name)
+    hair = cv2.imread(path, cv2.IMREAD_UNCHANGED)  # BGRA
+    if hair is None or hair.shape[2] != 4:
+        raise ValueError(f"Invalid hair asset: {name} (must be RGBA PNG)")
+    return hair
+
+def hair_reference_points(hair_bgra, filename):
+    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 (tune via 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)
+    return np.stack([pL, pR, pM], axis=0)
+
+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
+    hair_warp = cv2.warpAffine(hair_rgb, M, (W, H), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_TRANSPARENT)
+    a_warp = cv2.warpAffine(hair_a, M, (W, H), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_TRANSPARENT)
+    a = np.clip(a_warp * opacity, 0, 1)[..., None]
+    out = (a * hair_warp + (1 - a) * base_bgr).astype(np.uint8)
+    return out
+
+def apply_tryon(image, hairstyle, scale_pct, rot_deg, dx, dy, opacity):
+    if image is None:
+        return None, "Upload a photo or enable webcam."
+    if not hairstyle:
+        return np.array(image), "Pick a hairstyle first."
+
+    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)
+
+    dst = kpts.copy()
+    dst[:, 0] += dx
+    dst[:, 1] += dy
+
+    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)
+
+    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(img):
+    if img is None:
+        return None
+    p = os.path.join(tempfile.gettempdir(), "tryon_result.png")
+    Image.fromarray(img).save(p)
+    return p
+
+def hair_preview(hairstyle):
+    if not hairstyle:
+        return None
+    hair = load_hair_png(hairstyle)
+    h, w = hair.shape[:2]
+    tile = 16
+    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[:, :, 3:4] / 255.0)
+    comp = (a * hair[:, :, :3] + (1 - a) * bg_rgb).astype(np.uint8)
+    comp = cv2.cvtColor(comp, cv2.COLOR_BGR2RGB)
+    return comp
+
+# ---------------------- UI ----------------------
+def build_ui():
+    with gr.Blocks(title="Virtual Try-On (FR1–FR8)", css="""
+    .gradio-container {max-width: 980px; margin: auto;}
+    @media (max-width: 768px){ .gradio-container {padding: 8px;} }
+    """) as demo:
+        gr.Markdown("## Salon Hairstyle Virtual Try-On\nUpload or use webcam, pick a style from **Select Hairstyle**, adjust, then download.")
+
+        if not HAIR_FILES:
+            gr.Markdown("⚠️ **No hairstyle PNGs found.** Upload files into **`hair/`** (or `assets/hairstyles/`) and reload this Space.")
+
+        with gr.Tabs():
+            # -------- Photo (FR1,3–7) --------
+            with gr.Tab("Photo"):
+                with gr.Row():
+                    in_img = gr.Image(label="Upload photo (JPEG/PNG)", sources=["upload"], type="pil")
+                    hair = gr.Dropdown(
+                        choices=HAIR_FILES,
+                        value=(HAIR_FILES[0] if HAIR_FILES else None),
+                        label="Select Hairstyle (from 'hair/')",
+                        interactive=True
+                    )
+                with gr.Row():
+                    preview = gr.Image(label="Hairstyle Preview", height=260)
+                hair.change(fn=hair_preview, inputs=[hair], outputs=[preview])
+
+                with gr.Accordion("Alignment Controls", open=True):
+                    with gr.Row():
+                        scale = gr.Slider(50, 200, 100, 1, label="Scale %")
+                        rot = gr.Slider(-30, 30, 0, 1, label="Rotate (deg)")
+                    with gr.Row():
+                        dx = gr.Slider(-200, 200, 0, 1, label="Horizontal Nudge (px)")
+                        dy = gr.Slider(-200, 200, 0, 1, label="Vertical Nudge (px)")
+                    opacity = gr.Slider(0.2, 1.0, 1.0, 0.05, label="Hair Opacity")
+
+                out = gr.Image(label="Result Preview")
+                status = gr.Markdown()
+
+                run = gr.Button("Apply (Align & Overlay)")
+                run.click(
+                    fn=lambda im, h, s, r, dxv, dyv, op: apply_tryon(im, h, s, r, dxv, dyv, op),
+                    inputs=[in_img, hair, scale, rot, dx, dy, opacity],
+                    outputs=[out, status]
+                )
+
+                dl = gr.DownloadButton(label="Download Result", file_name="tryon.png")
+                dl.click(fn=save_png, inputs=[out], outputs=[dl])
+
+                gr.Markdown("Share this Space link after you make it public (FR-7).")
+
+            # -------- Webcam (FR2–FR6) --------
+            with gr.Tab("Webcam"):
+                cam = gr.Image(sources=["webcam"], streaming=True, type="pil", label="Enable camera to start")
+                hair2 = gr.Dropdown(choices=HAIR_FILES, value=(HAIR_FILES[0] if HAIR_FILES else None), label="Select Hairstyle")
+                scale2 = gr.Slider(50, 200, 100, 1, label="Scale %")
+                rot2   = gr.Slider(-25, 25, 0, 1, label="Rotate (deg)")
+                dx2    = gr.Slider(-150, 150, 0, 1, label="Horizontal Nudge (px)")
+                dy2    = gr.Slider(-150, 150, 0, 1, label="Vertical Nudge (px)")
+                opacity2 = gr.Slider(0.2, 1.0, 0.95, 0.05, label="Hair Opacity")
+                out2 = gr.Image(label="Live Preview")
+
+                def live(im, h, s, r, dxv, dyv, op):
+                    res, _ = apply_tryon(im, h, s, r, dxv, dyv, op)
+                    return res
+
+                cam.stream(live, inputs=[cam, hair2, scale2, rot2, dx2, dy2, opacity2], outputs=[out2])
+
+        return demo
+
+# Export for Spaces autostart
+app = build_ui()
+demo = app
+
+# Local dev
+if __name__ == "__main__":
+    app.launch()