Spaces:
Running
Running
Create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,26 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import gradio as gr
|
2 |
+
from pathlib import Path
|
3 |
+
from ui import build_ui
|
4 |
+
from overlay import apply_hairstyle
|
5 |
+
from segmentation import segment_image, estimate_landmarks
|
6 |
+
|
7 |
+
ASSETS_DIR = Path("assets/hairstyles")
|
8 |
+
|
9 |
+
def get_hairstyle_list():
|
10 |
+
return sorted([f.name for f in ASSETS_DIR.glob("*.png")])
|
11 |
+
|
12 |
+
def try_on(user_image, style_name):
|
13 |
+
if user_image is None or style_name is None:
|
14 |
+
return None
|
15 |
+
mask = segment_image(user_image)
|
16 |
+
landmarks = estimate_landmarks(user_image) # optional
|
17 |
+
style_path = str(ASSETS_DIR / style_name)
|
18 |
+
return apply_hairstyle(user_image, style_path, mask, landmarks)
|
19 |
+
|
20 |
+
def launch():
|
21 |
+
styles = get_hairstyle_list()
|
22 |
+
demo = build_ui(try_on, styles)
|
23 |
+
demo.launch()
|
24 |
+
|
25 |
+
if __name__ == "__main__":
|
26 |
+
launch()
|