app.py

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()