File size: 4,588 Bytes
79b274c e377d5f 79b274c e377d5f 79b274c e377d5f 79b274c e377d5f 79b274c e377d5f 79b274c e377d5f 79b274c e377d5f 79b274c e377d5f 79b274c e377d5f 79b274c e377d5f 79b274c e377d5f 79b274c e377d5f a5c4d00 e377d5f a5c4d00 79b274c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 |
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.Row():
with gr.Column(scale=2):
gr.Markdown("## 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.Column(scale=1):
gr.Markdown("## 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)
gr.Markdown("## Generate/Download Dataset")
generate_button = gr.Button("Generate Dataset")
generate_output = gr.File()
generate_button.click(generate_dataset, inputs=[dataset_name], outputs=generate_output)
gr.Markdown("## 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()
|