import gradio as gr import json import os import zipfile from datetime import datetime # Initialize dataset storage datasets = {} # Function to save dataset to JSONL file def save_dataset(dataset_name): if dataset_name in datasets: dataset = datasets[dataset_name] os.makedirs(f"{dataset_name}/images", exist_ok=True) with open(f"{dataset_name}/dataset.jsonl", "w") as f: for item in dataset: f.write(json.dumps(item) + "\n") # Function to load dataset from JSONL file def load_dataset(dataset_name, file): if file: os.makedirs(dataset_name, exist_ok=True) with open(f"{dataset_name}/dataset.jsonl", "r") as f: datasets[dataset_name] = [json.loads(line) for line in f] return f"Dataset '{dataset_name}' loaded successfully." return "No file uploaded." # Function to add image/prompt pair to dataset def add_to_dataset(dataset_name, image, prompt): if dataset_name not in datasets: datasets[dataset_name] = [] image_path = f"{dataset_name}/images/{datetime.now().timestamp()}.png" image.save(image_path) datasets[dataset_name].append({"image": image_path, "prompt": prompt}) return f"Added to dataset '{dataset_name}'." # Function to edit image/prompt pair in dataset def edit_dataset(dataset_name, index, new_prompt): if dataset_name in datasets and 0 <= index < len(datasets[dataset_name]): datasets[dataset_name][index]["prompt"] = new_prompt return f"Edited prompt at index {index} in dataset '{dataset_name}'." return "Invalid index." # Function to delete image/prompt pair from dataset def delete_from_dataset(dataset_name, index): if dataset_name in datasets and 0 <= index < len(datasets[dataset_name]): os.remove(datasets[dataset_name][index]["image"]) datasets[dataset_name].pop(index) return f"Deleted item at index {index} from dataset '{dataset_name}'." return "Invalid index." # Function to generate and download dataset as zip file def generate_dataset(dataset_name): if dataset_name in datasets: save_dataset(dataset_name) zip_filename = f"{dataset_name}.zip" with zipfile.ZipFile(zip_filename, 'w') as zipf: for root, _, files in os.walk(dataset_name): for file in files: zipf.write(os.path.join(root, file)) return zip_filename return "Dataset not found." # Gradio interface with gr.Blocks() as demo: gr.Markdown("# Dataset Builder") with gr.Tab("Create/Edit Dataset"): dataset_name = gr.Textbox(label="Dataset Name") image = gr.Image(label="Upload Image") prompt = gr.Textbox(label="Prompt") add_button = gr.Button("Add to Dataset") add_output = gr.Textbox() add_button.click(add_to_dataset, inputs=[dataset_name, image, prompt], outputs=add_output) index = gr.Number(label="Index to Edit/Delete", precision=0) new_prompt = gr.Textbox(label="New Prompt") edit_button = gr.Button("Edit Prompt") edit_output = gr.Textbox() edit_button.click(edit_dataset, inputs=[dataset_name, index, new_prompt], outputs=edit_output) delete_button = gr.Button("Delete from Dataset") delete_output = gr.Textbox() delete_button.click(delete_from_dataset, inputs=[dataset_name, index], outputs=delete_output) with gr.Tab("Upload Dataset"): upload_dataset_name = gr.Textbox(label="Dataset Name") upload_file = gr.File(label="Upload JSONL File") upload_button = gr.Button("Upload Dataset") upload_output = gr.Textbox() upload_button.click(load_dataset, inputs=[upload_dataset_name, upload_file], outputs=upload_output) with gr.Tab("Generate/Download Dataset"): generate_button = gr.Button("Generate Dataset") generate_output = gr.File() generate_button.click(generate_dataset, inputs=[dataset_name], outputs=generate_output) with gr.Tab("View Dataset"): view_dataset_name = gr.Textbox(label="Dataset Name") view_button = gr.Button("View Dataset") view_output = gr.JSON() view_button.click(lambda name: datasets.get(name, []), inputs=[view_dataset_name], outputs=view_output) demo.launch()