Update app.py

app.py

@@ -4,14 +4,25 @@ from ui import build_ui
 from overlay import apply_hairstyle
 from segmentation import segment_image, estimate_landmarks
 
+# Folder containing your PNG hairstyle overlays
 ASSETS_DIR = Path("assets/hairstyles")
 
 def get_hairstyle_list():
-    return sorted([f.name for f in ASSETS_DIR.glob("*.png")])
+    """Return sorted list of PNG files in the assets/hairstyles folder."""
+    if not ASSETS_DIR.exists():
+        print("⚠️ hairstyles folder not found:", ASSETS_DIR)
+        return []
+    styles = sorted([f.name for f in ASSETS_DIR.glob("*.png") if f.is_file()])
+    print("✅ Found hairstyles:", styles)
+    return styles
 
 def try_on(user_image, style_name):
-    if user_image is None or style_name is None:
+    """Main try-on function: segments hair, aligns style, blends result."""
+    if user_image is None:
         return None
+    if not style_name:
+        print("⚠️ No hairstyle selected")
+        return user_image
     mask = segment_image(user_image)
     landmarks = estimate_landmarks(user_image)  # optional
     style_path = str(ASSETS_DIR / style_name)
@@ -19,6 +30,8 @@ def try_on(user_image, style_name):
 
 def launch():
     styles = get_hairstyle_list()
+    if not styles:
+        print("⚠️ No PNG hairstyles found in", ASSETS_DIR)
     demo = build_ui(try_on, styles)
     demo.launch()