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