import gradio as gr
import cv2
import numpy as np
from Project.aligned_image.aligned_images import align
from Project.scripts.inference import inference
from Project.notebook.out import out

def final(image, style):
    myimg = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    cv2.imwrite('static/img_in/in.jpg', myimg)
    align()
    inference()
    out(style)
    result = cv2.imread('static/img_out/inference_results/00000.jpg')
    aligned = cv2.imread('static/img_aligned/in_01.png')
    return cv2.cvtColor(result, cv2.COLOR_RGB2BGR), cv2.cvtColor(aligned, cv2.COLOR_RGB2BGR)

# 创建 Gradio 接口
style_options = {
    "Emotion": {'Angry': 'angry', 'Surprised': 'surprised'},
    "Celebrity": {'Beyonce': 'Beyonce', 'Hilary Clinton': 'Hilary_clinton', 'Johnny Depp': 'Jhonny Depp', 'Taylor Swift': 'Taylor Swift', 'Trump': 'trump'},
    "Hair Style": {'Afro': 'afro', 'Bowlcut': 'bowlcut', 'Curly Hair': 'curly hair', 'Purple Hair': 'purple hair'}
}

# 定义更新风格选项的函数
def update_styles(style_type):
    if not style_type:
        return gr.Dropdown(choices=[])
    return gr.Dropdown(choices=list(style_options[style_type].keys()))

# 创建 Gradio 界面
with gr.Blocks() as demo:
    gr.Markdown("# Image Style Transfer")
    gr.Markdown("### This app is based on styleclip. Choose a style type and a style from the dropdowns below.")
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="numpy", label="Upload Image")
            style_type_dropdown = gr.Dropdown(choices=list(style_options.keys()), label="Style Type")
            style_dropdown = gr.Dropdown(choices=["Angry", "Curly Hair", "Taylor Swift"], label="Style")
            style_type_dropdown.change(fn=update_styles, inputs=style_type_dropdown, outputs=style_dropdown)
            with gr.Row():
                clear_button = gr.Button("Clear")
                submit_button = gr.Button("Submit")
        with gr.Column():
            output_image = gr.Image(type="numpy", label="Result Image")
            aligned_image = gr.Image(type="numpy", label="Aligned Image")
    
    def on_submit(image, style_type, style):
        style_value = style_options[style_type][style]
        return final(image, style_value)
    
    def on_clear():
        return None, None, None, None, None
    
    clear_button.click(fn=on_clear, inputs=[], outputs=[image_input, style_type_dropdown, style_dropdown, output_image, aligned_image])
    submit_button.click(fn=on_submit, inputs=[image_input, style_type_dropdown, style_dropdown], outputs=[output_image, aligned_image])

    examples = gr.Examples(
        examples=[
            ["static/img/example1.jpg", "Emotion", "Angry"],
            ["static/img/example2.jpg", "Celebrity", "Taylor Swift"],
            ["static/img/example3.jpg", "Hair Style", "Curly Hair"],
        ],
        inputs=[image_input, style_type_dropdown, style_dropdown],
    )
# 启动 Gradio 应用
demo.launch()