|
import gradio as gr |
|
import os |
|
import csv |
|
|
|
root = 'saved' |
|
prompt_path = 'assets/ViLG-300.csv' |
|
|
|
|
|
def load_prompts(path): |
|
if os.path.basename(path) == 'ViLG-300.csv': |
|
def csv_to_dict(file_path): |
|
result_dict = {} |
|
with open(file_path, 'r', encoding='utf-8') as csv_file: |
|
csv_reader = csv.DictReader(csv_file, delimiter=',') |
|
for row in csv_reader: |
|
prompt = row['\ufeffPrompt'] |
|
text = row['文本'] |
|
category = row['类别'] |
|
source = row['来源'] |
|
result_dict[prompt] = {'prompt': prompt, 'text': text, 'category': category, 'source': source} |
|
return result_dict |
|
data = list(csv_to_dict(path).keys()) |
|
else: |
|
return NotImplementedError |
|
return data |
|
|
|
|
|
prompts = load_prompts(prompt_path) |
|
|
|
|
|
def load_images(methods, idx): |
|
idx = int(idx) |
|
prompt = prompts[idx].strip() |
|
images = [] |
|
for method in methods: |
|
image = os.path.join(root, method, f'{idx}.jpg') |
|
images.append((image, method)) |
|
return prompt, images |
|
|
|
|
|
def load_methods(): |
|
methods = os.listdir(root) |
|
return methods |
|
|
|
|
|
def main(): |
|
with gr.Blocks() as demo: |
|
gr.Markdown("# Text to Image Models Comparison") |
|
with gr.Row(): |
|
idx = gr.Number(value=0, label='Index') |
|
prompt = gr.Textbox() |
|
methods = gr.Dropdown(multiselect=True, choices=load_methods(), value=load_methods(), label='Methods') |
|
gallery = gr.Gallery(show_label=False, object_fit='fill', height=600, columns=5) |
|
idx.change(load_images, [methods, idx], [prompt, gallery]) |
|
methods.change(load_images, [methods, idx], [prompt, gallery]) |
|
demo.launch() |
|
|
|
|
|
if __name__ == '__main__': |
|
main() |
|
|