Update app.py

app.py

@@ -5,7 +5,7 @@ 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, "hair"),                    # <- your folder
     os.path.join(BASE_DIR, "assets", "hairstyles"),
     os.path.join(BASE_DIR, "assets", "Hairstyles"),
     os.path.join(BASE_DIR, "hairstyles"),
@@ -15,28 +15,24 @@ 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
+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
 
-# ---------------------- Dependencies ----------------------
+# ---------------------- MediaPipe ----------------------
 try:
-    import mediapipe as mp  # FaceMesh + SelfieSeg
+    import mediapipe as mp
 except Exception as e:
-    raise RuntimeError(
-        f"Failed to import mediapipe. Check requirements.txt pins. Details: {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():
-    """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:
@@ -55,11 +51,9 @@ def load_meta():
         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,
@@ -73,124 +67,4 @@ def detect_face_keypoints(img_bgr):
         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} (must be RGBA PNG)")
-    return 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]
-    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_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)
-
-    # User nudges on destination points
-    dst = kpts.copy()
-    dst[:, 0] += dx
-    dst[:, 1] += dy
-
-    # 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
-
-    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(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():
-                    i