Create app.py

app.py

@@ -0,0 +1,98 @@
+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()