File size: 3,757 Bytes
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 |
import gradio as gr
import os
import json
import zipfile
from datasets import Dataset, DatasetDict
# Directory to store datasets
DATASETS_DIR = "datasets"
# Ensure the datasets directory exists
os.makedirs(DATASETS_DIR, exist_ok=True)
# Function to save dataset to persistent storage
def save_dataset(dataset_name, data):
dataset_path = os.path.join(DATASETS_DIR, f"{dataset_name}.jsonl")
with open(dataset_path, 'w') as f:
for item in data:
f.write(json.dumps(item) + "\n")
# Function to load dataset from persistent storage
def load_dataset(dataset_name):
dataset_path = os.path.join(DATASETS_DIR, f"{dataset_name}.jsonl")
if os.path.exists(dataset_path):
with open(dataset_path, 'r') as f:
data = [json.loads(line) for line in f]
return data
return []
# Function to generate and download dataset as a zip file
def generate_dataset(dataset_name):
dataset_path = os.path.join(DATASETS_DIR, f"{dataset_name}.jsonl")
if not os.path.exists(dataset_path):
return None
# Create a zip file
zip_filename = f"{dataset_name}.zip"
with zipfile.ZipFile(zip_filename, 'w') as zipf:
# Add images and JSONL file to the zip
with open(dataset_path, 'r') as f:
data = [json.loads(line) for line in f]
jsonl_path = os.path.join(DATASETS_DIR, f"{dataset_name}.jsonl")
zipf.write(jsonl_path, arcname=f"{dataset_name}.jsonl")
for item in data:
image_path = item['image']
zipf.write(image_path, arcname=os.path.join('images', os.path.basename(image_path)))
return zip_filename
# Gradio interface functions
def add_image_prompt(dataset_name, image, prompt):
data = load_dataset(dataset_name)
image_path = os.path.join(DATASETS_DIR, os.path.basename(image.name))
image.save(image_path)
data.append({"image": image_path, "prompt": prompt})
save_dataset(dataset_name, data)
return data
def edit_image_prompt(dataset_name, index, new_prompt):
data = load_dataset(dataset_name)
if 0 <= index < len(data):
data[index]['prompt'] = new_prompt
save_dataset(dataset_name, data)
return data
def delete_image_prompt(dataset_name, index):
data = load_dataset(dataset_name)
if 0 <= index < len(data):
del data[index]
save_dataset(dataset_name, data)
return data
def get_datasets():
return [f.split('.')[0] for f in os.listdir(DATASETS_DIR) if f.endswith('.jsonl')]
# Gradio UI
with gr.Blocks() as demo:
dataset_name = gr.Textbox(label="Dataset Name")
image = gr.Image(label="Upload Image")
prompt = gr.Textbox(label="Prompt")
add_button = gr.Button("Add Image/Prompt")
edit_index = gr.Number(label="Edit Index", value=0)
new_prompt = gr.Textbox(label="New Prompt")
edit_button = gr.Button("Edit Prompt")
delete_index = gr.Number(label="Delete Index", value=0)
delete_button = gr.Button("Delete Image/Prompt")
generate_button = gr.Button("Generate Dataset")
download_button = gr.Button("Download Dataset")
dataset_list = gr.Dropdown(label="Select Dataset", choices=get_datasets())
dataset_display = gr.Dataframe(headers=["Image", "Prompt"])
add_button.click(add_image_prompt, [dataset_name, image, prompt], dataset_display)
edit_button.click(edit_image_prompt, [dataset_name, edit_index, new_prompt], dataset_display)
delete_button.click(delete_image_prompt, [dataset_name, delete_index], dataset_display)
generate_button.click(generate_dataset, [dataset_name], None)
download_button.click(generate_dataset, [dataset_name], None, file_name=f"{dataset_name}.zip")
dataset_list.change(load_dataset, [dataset_list], dataset_display)
demo.launch()
|