|
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() |
|
|