Update app.py

app.py

@@ -1,98 +1,98 @@
 import gradio as gr
-import os
 import json
+import os
 import zipfile
-from datasets import Dataset, DatasetDict
-
-# Directory to store datasets
-DATASETS_DIR = "datasets"
+from datetime import datetime
 
-# Ensure the datasets directory exists
-os.makedirs(DATASETS_DIR, exist_ok=True)
+# Initialize dataset storage
+datasets = {}
 
-# 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 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 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
+# 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."
 
-    # 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))
+# 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)
-    data.append({"image": image_path, "prompt": prompt})
-    save_dataset(dataset_name, data)
-    return data
+    datasets[dataset_name].append({"image": image_path, "prompt": prompt})
+    return f"Added to dataset '{dataset_name}'."
 
-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
+# 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."
 
-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
+# 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."
 
-def get_datasets():
-    return [f.split('.')[0] for f in os.listdir(DATASETS_DIR) if f.endswith('.jsonl')]
+# 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 UI
+# Gradio interface
 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)
+    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_button.click(add_to_dataset, inputs=[dataset_name, image, prompt], outputs="text")
+        
+        index = gr.Number(label="Index to Edit/Delete", precision=0)
+        new_prompt = gr.Textbox(label="New Prompt")
+        edit_button = gr.Button("Edit Prompt")
+        edit_button.click(edit_dataset, inputs=[dataset_name, index, new_prompt], outputs="text")
+        
+        delete_button = gr.Button("Delete from Dataset")
+        delete_button.click(delete_from_dataset, inputs=[dataset_name, index], outputs="text")
+    
+    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_button.click(load_dataset, inputs=[upload_dataset_name, upload_file], outputs="text")
+    
+    with gr.Tab("Generate/Download Dataset"):
+        generate_button = gr.Button("Generate Dataset")
+        generate_button.click(generate_dataset, inputs=[dataset_name], outputs="file")
+    
+    with gr.Tab("View Dataset"):
+        view_dataset_name = gr.Textbox(label="Dataset Name")
+        view_button = gr.Button("View Dataset")
+        view_button.click(lambda name: datasets.get(name, []), inputs=[view_dataset_name], outputs="json")
 
 demo.launch()