Update app.py

app.py

@@ -1,187 +1,196 @@
-import os
-import cv2
-import numpy as np
-import gradio as gr
-import mediapipe as mp
+import os, json, tempfile
+import cv2, numpy as np, gradio as gr
+from PIL import Image
+
+# ---------------------- Paths (hair/ first) ----------------------
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+CANDIDATES = [
+    os.path.join(BASE_DIR, "hair"),                    # <- your current folder
+    os.path.join(BASE_DIR, "assets", "hairstyles"),
+    os.path.join(BASE_DIR, "assets", "Hairstyles"),
+    os.path.join(BASE_DIR, "hairstyles"),
+]
+HAIR_DIR = None
+for p in CANDIDATES:
+    if os.path.isdir(p):
+        HAIR_DIR = p
+        break
+if HAIR_DIR is None:  # create the 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
+
+# ---------------------- Dependencies ----------------------
+try:
+    import mediapipe as mp  # FaceMesh + SelfieSeg
+except Exception as e:
+    raise RuntimeError(
+        f"Failed to import mediapipe. Check requirements.txt pins. Details: {e}"
+    )
 
-HAIR_DIR = "assets/hairstyles"
-
-# ---------- MediaPipe FaceMesh setup ----------
 mp_face_mesh = mp.solutions.face_mesh
-FACE_LANDMARKS = {
-    "left_eye_outer": 33,   # standard mesh indices
-    "right_eye_outer": 263,
-    "mid_forehead": 10      # glabella/forehead area (approx)
-}
+mp_selfie_seg = mp.solutions.selfie_segmentation
+
+LM = {"left_eye_outer": 33, "right_eye_outer": 263, "mid_forehead": 10}
 
+# ---------------------- Helpers ----------------------
 def load_hairstyles():
-    files = [f for f in os.listdir(HAIR_DIR) if f.lower().endswith(".png")]
+    """Return sorted list of .png files in HAIR_DIR."""
+    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 detect_keypoints(image_bgr):
-    """Returns 3 keypoints (x,y) in image coords (float)."""
-    h, w = image_bgr.shape[:2]
+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):
+    """Return 3 keypoints (left eye outer, right eye outer, mid-forehead) or None."""
+    h, w = img_bgr.shape[:2]
     with mp_face_mesh.FaceMesh(
-        static_image_mode=True,
-        max_num_faces=1,
-        refine_landmarks=True,
+        static_image_mode=True, max_num_faces=1, refine_landmarks=True,
         min_detection_confidence=0.6
-    ) as face_mesh:
-        img_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
-        res = face_mesh.process(img_rgb)
+    ) 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 to_xy(idx):
-            return np.array([lm[idx].x * w, lm[idx].y * h], dtype=np.float32)
-        p_left  = to_xy(FACE_LANDMARKS["left_eye_outer"])
-        p_right = to_xy(FACE_LANDMARKS["right_eye_outer"])
-        p_mid   = to_xy(FACE_LANDMARKS["mid_forehead"])
-        return np.stack([p_left, p_right, p_mid], axis=0)
+        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):
+    """Rough head isolation using selfie segmentation + feathering."""
+    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}")
+        raise ValueError(f"Invalid hair asset: {name} (must be RGBA PNG)")
     return hair
 
-def hair_reference_points(hair_bgra):
-    """
-    Define 3 reference anchor points on the hair image (in hair coords).
-    Tune once per style so that these points roughly correspond to
-    left eye outer, right eye outer, and mid-forehead anchors.
-    For a simple default, place anchors across the lower edge of hair.
-    """
+def hair_reference_points(hair_bgra, filename):
+    """Three anchors on hair image; override via meta.json if present."""
     h, w = hair_bgra.shape[:2]
-    # Default guesses; adjust per asset for better alignment.
+    if filename in META:
+        pts = np.array(META[filename], dtype=np.float32)
+        if pts.shape == (3, 2):
+            return pts
+    # Defaults (works for many styles; refine 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_blend(base_bgr, hair_bgra, M, opacity=1.0):
-    h_img, w_img = base_bgr.shape[:2]
-    # Split hair into BGR + A
-    hair_bgr = hair_bgra[:, :, :3]
-    hair_a   = hair_bgra[:, :, 3] / 255.0
-
-    # Warp both color and alpha with same affine
-    hair_warp = cv2.warpAffine(hair_bgr, M, (w_img, h_img), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_TRANSPARENT)
-    alpha_warp = cv2.warpAffine(hair_a, M, (w_img, h_img), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_TRANSPARENT)
-
-    alpha = np.clip(alpha_warp * opacity, 0, 1)[..., None]  # HxWx1
-    out = (alpha * hair_warp + (1 - alpha) * base_bgr).astype(np.uint8)
+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 tryon(
-    image, 
-    hairstyle, 
-    scale_pct=100, 
-    rot_deg=0, 
-    dx=0, dy=0, 
-    opacity=1.0,
-    mode="Photo"
-):
+def apply_tryon(image, hairstyle, scale_pct, rot_deg, dx, dy, opacity):
     if image is None:
-        return None, "Please provide an image or turn on webcam."
+        return None, "Upload a photo or enable webcam."
+    if not hairstyle:
+        return np.array(image), "Pick a hairstyle first."
 
-    # Convert PIL->BGR
     img_bgr = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
-    keypts = detect_keypoints(img_bgr)
-    if keypts is None:
-        return image, "No face detected. Try a clearer, front-facing photo."
+
+    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)
+    hair_pts = hair_reference_points(hair, hairstyle)
 
-    # Compute affine (hair -> face) using three points
-    dst_pts = keypts.copy()
-    # Apply manual user tweaks on destination:
-    dst_pts[:, 0] += dx
-    dst_pts[:, 1] += dy
+    # User nudges on destination points
+    dst = kpts.copy()
+    dst[:, 0] += dx
+    dst[:, 1] += dy
 
-    # Scale and rotate around hair center before solving affine
+    # Scale + rotate hair anchors around their 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
 
-    # Solve affine transform
-    M, _ = cv2.estimateAffinePartial2D(hair_pts_adj, dst_pts, method=cv2.LMEDS)
-
+    M, _ = cv2.estimateAffinePartial2D(hair_pts_adj, dst, method=cv2.LMEDS)
     if M is None:
-        return image, "Could not compute alignment. Try another image/hairstyle."
+        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_bgr = warp_and_blend(img_bgr, hair, M, opacity=opacity)
-    out_rgb = cv2.cvtColor(out_bgr, cv2.COLOR_BGR2RGB)
+    out_rgb = cv2.cvtColor(out, cv2.COLOR_BGR2RGB)
     return out_rgb, "OK"
 
-def build_ui():
-    with gr.Blocks(title="Salon Hairstyle Virtual Try-On") as demo:
-        gr.Markdown("# Salon Hairstyle Virtual Try-On\nTry different hairstyles in real-time or on a photo.")
+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
+    # Show the raw PNG on checkerboard background for visibility
+    hair = load_hair_png(hairstyle)
+    h, w = hair.shape[:2]
+    # Make checkerboard
+    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 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 Tab (FR1,3–7) ----------------
             with gr.Tab("Photo"):
                 with gr.Row():
-                    in_img = gr.Image(type="pil", label="Upload photo", sources=["upload"], height=420)
-                    gallery = gr.Radio(choices=HAIR_FILES, value=HAIR_FILES[0] if HAIR_FILES else None, label="Hairstyles")
-                with gr.Accordion("Alignment & Style Controls", open=True):
-                    with gr.Row():
-                        scale = gr.Slider(50, 200, value=100, step=1, label="Scale %")
-                        rot   = gr.Slider(-30, 30, value=0, step=1, label="Rotate (deg)")
-                    with gr.Row():
-                        dx    = gr.Slider(-200, 200, value=0, step=1, label="Horizontal Nudge (px)")
-                        dy    = gr.Slider(-200, 200, value=0, step=1, label="Vertical Nudge (px)")
-                    opacity = gr.Slider(0.2, 1.0, value=1.0, step=0.05, label="Hair Opacity")
-
-                out_img = gr.Image(label="Result", height=420)
-                status  = gr.Markdown()
-
-                run_btn = gr.Button("Apply")
-                save_btn = gr.Button("Save Result")
-
-                def on_save(img):
-                    # Gradio will let users right-click or use built-in download;
-                    # Optionally, return img so it can be saved from gallery.
-                    return img
-
-                run_btn.click(
-                    fn=lambda im, h, s, r, dxv, dyv, op: tryon(im, h, s, r, dxv, dyv, op, mode="Photo"),
-                    inputs=[in_img, gallery, scale, rot, dx, dy, opacity],
-                    outputs=[out_img, status]
-                )
-                save_btn.click(fn=on_save, inputs=[out_img], outputs=[out_img])
-
-            with gr.Tab("Webcam (Live)"):
-                cam = gr.Image(sources=["webcam"], streaming=True, type="pil", label="Webcam")
-                hair2 = gr.Radio(choices=HAIR_FILES, value=HAIR_FILES[0] if HAIR_FILES else None, label="Hairstyles")
-                scale2 = gr.Slider(50, 200, value=100, step=1, label="Scale %")
-                rot2   = gr.Slider(-25, 25, value=0, step=1, label="Rotate (deg)")
-                dx2    = gr.Slider(-150, 150, value=0, step=1, label="Horizontal Nudge (px)")
-                dy2    = gr.Slider(-150, 150, value=0, step=1, label="Vertical Nudge (px)")
-                opacity2 = gr.Slider(0.2, 1.0, value=0.95, step=0.05, label="Hair Opacity")
-
-                out2 = gr.Image(label="Live Result")
-
-                def live_fn(im, h, s, r, dxv, dyv, op):
-                    res, _ = tryon(im, h, s, r, dxv, dyv, op, mode="Webcam")
-                    return res
-
-                cam.stream(
-                    fn=live_fn,
-                    inputs=[cam, hair2, scale2, rot2, dx2, dy2, opacity2],
-                    outputs=[out2],
-                    time_limit=0.0 # continuous
-                )
-        return demo
-
-demo = build_ui()
-
-if __name__ == "__main__":
-    demo.launch()
+                    i