import gradio as gr
from pathlib import Path
from ui import build_ui
from overlay import apply_hairstyle
from segmentation import segment_image, estimate_landmarks

ASSETS_DIR = Path("assets/hairstyles")

def get_hairstyle_list():
    return sorted([f.name for f in ASSETS_DIR.glob("*.png")])

def try_on(user_image, style_name):
    if user_image is None or style_name is None:
        return None
    mask = segment_image(user_image)
    landmarks = estimate_landmarks(user_image)  # optional
    style_path = str(ASSETS_DIR / style_name)
    return apply_hairstyle(user_image, style_path, mask, landmarks)

def launch():
    styles = get_hairstyle_list()
    demo = build_ui(try_on, styles)
    demo.launch()

if __name__ == "__main__":
    launch()