import gradio as gr
from feature_combiner import FeatureCombiner
from pipeline import Pipeline_demo
# 数据集选择字典
DATASETS = {
    "angiography.CHUAC": "CHUAC dataset",
    "angiography.DCA1": "DCA1 dataset",
    "angiography.XCAD": "XCAD dataset",
    "retina.DRIVE": "DRIVE dataset",
    "retina.CHASEDB1": "CHASEDB1 dataset",
    "crack.Crack500": "Crack500 dataset"
}

JSON_FILES = {
    "crack":"dataset_curvilinear/crack/crack.json",
    "angiography":"dataset_curvilinear/angiography/angiography.json",
    "retina":"dataset_curvilinear/retina/retina.json"
}

is_cpu= True
pipeline_demo = Pipeline_demo(is_cpu=is_cpu) 

def generate_prompts(dataset):
    """
    根据选择的数据集生成prompts。
    返回semantic map prompt和image prompt。
    """

    json_index = dataset.split(".")[0]
    json_file = JSON_FILES[json_index]
    fc = FeatureCombiner(json_file, DATASETS[dataset])
    caption_seg, caption_img = fc.combine()
    return caption_seg, caption_img

def generate_semantic_map(semantic_map_prompt):
    """
    根据semantic map prompt生成semantic map。
    返回semantic map的图像。
    """
    # 示例函数体，你需要填充实际的图像生成逻辑
    return pipeline_demo.generate_semantic_map(semantic_map_prompt)


def generate_image(image_prompt,semantic_map_image):
    """
    根据image prompt生成图像。
    返回生成的图像。
    """
    # 示例函数体，你需要填充实际的图像生成逻辑
    return pipeline_demo.generate_image(image_prompt, semantic_map_image)

with gr.Blocks() as demo:
    gr.Markdown("# DEMO for Expanding Curvilinear Object Segmentation datasets")
    with gr.Row():
        with gr.Column(scale=1):
            dataset_dropdown = gr.Dropdown(label="Select Dataset", choices=list(DATASETS.keys()),value="crack.Crack500")
            generate_prompts_button = gr.Button("1. Generate Prompts")
            gr.Markdown("""
            #### Process:
            1. Generate Prompts
            2. Generate Semantic Map based on the prompt
            3. Generate Image based on the prompt and Semantic Map  
            NOTE: The time it takes for the CPU provided by Huggingface to complete the generation is too long.  It is recommended that users modify is_cpu=False in app.py to deploy on their own devices.
            """)
        with gr.Column(scale=2):
            semantic_map_prompt_text = gr.Textbox(label="Semantic Map Prompt", lines=5, placeholder="Semantic map prompt will be shown here...", interactive=True)
            image_prompt_text = gr.Textbox(label="Image Prompt", lines=5, placeholder="Image prompt will be shown here...", interactive=True)

    with gr.Row():
        with gr.Column(scale=1):
            generate_semantic_map_button = gr.Button("2. Generate Semantic Map")
            semantic_map_image = gr.Image(width=512, height=512, label="Semantic Map",interactive=False)
        with gr.Column(scale=1):
            generate_image_button = gr.Button("3. Generate Image")
            generated_image = gr.Image(width=512, height=512, label="Generated Image")

    # 绑定函数
    generate_prompts_button.click(generate_prompts, inputs=[dataset_dropdown], outputs=[semantic_map_prompt_text, image_prompt_text])
    generate_semantic_map_button.click(generate_semantic_map, inputs=[semantic_map_prompt_text], outputs=[semantic_map_image])
    generate_image_button.click(generate_image, inputs=[image_prompt_text,semantic_map_image], outputs=[generated_image])

demo.launch()