|
|
|
|
|
from __future__ import annotations |
|
|
|
import pathlib |
|
|
|
import gradio as gr |
|
|
|
from model import Model |
|
|
|
DESCRIPTION = """# [HairCLIP](https://github.com/wty-ustc/HairCLIP) |
|
|
|
<center><img id="teaser" src="https://raw.githubusercontent.com/wty-ustc/HairCLIP/main/assets/teaser.png" alt="teaser"></center> |
|
""" |
|
|
|
|
|
def load_hairstyle_list() -> list[str]: |
|
with open("HairCLIP/mapper/hairstyle_list.txt") as f: |
|
lines = [line.strip() for line in f.readlines()] |
|
lines = [line[:-10] for line in lines] |
|
return lines |
|
|
|
|
|
def set_example_image(example: list) -> dict: |
|
return gr.Image(value=example[0]) |
|
|
|
|
|
def update_step2_components(choice: str) -> tuple[dict, dict]: |
|
return ( |
|
gr.Dropdown(visible=choice in ["hairstyle", "both"]), |
|
gr.Textbox(visible=choice in ["color", "both"]), |
|
) |
|
|
|
|
|
model = Model() |
|
|
|
with gr.Blocks(css="style.css") as demo: |
|
gr.Markdown(DESCRIPTION) |
|
with gr.Group(): |
|
gr.Markdown("## Step 1") |
|
with gr.Row(): |
|
with gr.Column(): |
|
with gr.Row(): |
|
input_image = gr.Image(label="Input Image", type="filepath") |
|
with gr.Row(): |
|
preprocess_button = gr.Button("Preprocess") |
|
with gr.Column(): |
|
aligned_face = gr.Image(label="Aligned Face", type="pil", interactive=False) |
|
with gr.Column(): |
|
reconstructed_face = gr.Image(label="Reconstructed Face", type="numpy") |
|
latent = gr.State() |
|
|
|
with gr.Row(): |
|
paths = sorted(pathlib.Path("images").glob("*.jpg")) |
|
gr.Examples(examples=[[path.as_posix()] for path in paths], inputs=input_image) |
|
|
|
with gr.Group(): |
|
gr.Markdown("## Step 2") |
|
with gr.Row(): |
|
with gr.Column(): |
|
with gr.Row(): |
|
editing_type = gr.Radio( |
|
label="Editing Type", choices=["hairstyle", "color", "both"], value="both", type="value" |
|
) |
|
with gr.Row(): |
|
hairstyles = load_hairstyle_list() |
|
hairstyle_index = gr.Dropdown(label="Hairstyle", choices=hairstyles, value="afro", type="index") |
|
with gr.Row(): |
|
color_description = gr.Textbox(label="Color", value="red") |
|
with gr.Row(): |
|
run_button = gr.Button("Run") |
|
|
|
with gr.Column(): |
|
result = gr.Image(label="Result") |
|
|
|
preprocess_button.click(fn=model.detect_and_align_face, inputs=input_image, outputs=aligned_face) |
|
aligned_face.change(fn=model.reconstruct_face, inputs=aligned_face, outputs=[reconstructed_face, latent]) |
|
editing_type.change(fn=update_step2_components, inputs=editing_type, outputs=[hairstyle_index, color_description]) |
|
run_button.click( |
|
fn=model.generate, |
|
inputs=[ |
|
editing_type, |
|
hairstyle_index, |
|
color_description, |
|
latent, |
|
], |
|
outputs=result, |
|
) |
|
|
|
if __name__ == "__main__": |
|
demo.queue(max_size=10).launch() |
|
|
