Update app.py

app.py

@@ -1,107 +1,243 @@
 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
+import json
+from io import BytesIO
+import base64
+from PIL import Image
+import uuid
+import tempfile
+
+def save_dataset_to_zip(dataset_name, dataset):
+    # Create a temporary directory
+    temp_dir = tempfile.mkdtemp()
+    dataset_path = os.path.join(temp_dir, dataset_name)
+    os.makedirs(dataset_path, exist_ok=True)
+    images_dir = os.path.join(dataset_path, 'images')
+    os.makedirs(images_dir, exist_ok=True)
+    annotations = []
+    for idx, entry in enumerate(dataset):
+        image_data = entry['image']
+        prompt = entry['prompt']
+        # Save image to images directory
+        image_filename = f"{uuid.uuid4().hex}.png"
+        image_path = os.path.join(images_dir, image_filename)
+        image = Image.open(BytesIO(base64.b64decode(image_data.split(",")[1])))
+        image.save(image_path)
+        # Add annotation
+        annotations.append({
+            'file_name': os.path.join('images', image_filename),
+            'text': prompt
+        })
+    # Save annotations to JSONL file
+    annotations_path = os.path.join(dataset_path, 'annotations.jsonl')
+    with open(annotations_path, 'w') as f:
+        for ann in annotations:
+            f.write(json.dumps(ann) + '\n')
+    # Create a zip file
+    zip_buffer = BytesIO()
+    with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zipf:
+        for root, dirs, files in os.walk(dataset_path):
+            for file in files:
+                abs_file = os.path.join(root, file)
+                rel_file = os.path.relpath(abs_file, dataset_path)
+                zipf.write(abs_file, rel_file)
+    zip_buffer.seek(0)
+    return zip_buffer
+
+def load_dataset_from_zip(zip_file):
+    temp_dir = tempfile.mkdtemp()
+    with zipfile.ZipFile(zip_file.name, 'r') as zip_ref:
+        zip_ref.extractall(temp_dir)
+    # Assuming the dataset folder is the first folder in the zip
+    dataset_name = os.listdir(temp_dir)[0]
+    dataset_path = os.path.join(temp_dir, dataset_name)
+    dataset = []
+    images_dir = os.path.join(dataset_path, 'images')
+    annotations_path = os.path.join(dataset_path, 'annotations.jsonl')
+    if os.path.exists(annotations_path):
+        with open(annotations_path, 'r') as f:
+            for line in f:
+                ann = json.loads(line)
+                file_name = ann['file_name']
+                prompt = ann['text']
+                image_path = os.path.join(dataset_path, file_name)
+                # Read image and convert to base64
+                with open(image_path, 'rb') as img_f:
+                    image_bytes = img_f.read()
+                    encoded = base64.b64encode(image_bytes).decode()
+                    mime_type = "image/png"
+                    image_data = f"data:{mime_type};base64,{encoded}"
+                dataset.append({
+                    'image': image_data,
+                    'prompt': prompt
+                })
+    return dataset_name, dataset
+
+def display_dataset_html(dataset):
+    if dataset:
+        html_content = ""
+        for idx, entry in enumerate(dataset):
+            image_data = entry['image']
+            prompt = entry['prompt']
+            html_content += f"""
+            <div style="display: flex; align-items: center; margin-bottom: 10px;">
+                <div style="width: 50px;">{idx}</div>
+                <img src="{image_data}" alt="Image {idx}" style="max-height: 100px; margin-right: 10px;"/>
+                <div>{prompt}</div>
+            </div>
+            """
+        return html_content
+    else:
+        return "<div>No entries in dataset.</div>"
+
 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()
+    gr.Markdown("<h1 style='text-align: center; margin-bottom: 20px;'>Dataset Builder</h1>")
+    datasets = gr.State({})
+    current_dataset_name = gr.State("")
+    dataset_selector = gr.Dropdown(label="Select Dataset", interactive=True)
+    dataset_html = gr.HTML()
+    message_box = gr.Textbox(interactive=False, label="Message")
+
+    with gr.Tab("Create / Upload Dataset"):
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("### Create a New Dataset")
+                dataset_name_input = gr.Textbox(label="New Dataset Name")
+                create_button = gr.Button("Create Dataset")
+            with gr.Column():
+                gr.Markdown("### Upload Existing Dataset")
+                upload_input = gr.File(label="Upload Dataset Zip", file_types=['.zip'])
+                upload_button = gr.Button("Upload Dataset")
+
+    def create_dataset(name, datasets):
+        if not name:
+            return gr.update(), "Please enter a dataset name."
+        if name in datasets:
+            return gr.update(), f"Dataset '{name}' already exists."
+        datasets[name] = []
+        return gr.Dropdown.update(choices=list(datasets.keys()), value=name), f"Dataset '{name}' created."
+
+    create_button.click(
+        create_dataset,
+        inputs=[dataset_name_input, datasets],
+        outputs=[dataset_selector, message_box]
+    )
+
+    def upload_dataset(zip_file, datasets):
+        if zip_file is None:
+            return gr.update(), "Please upload a zip file."
+        dataset_name, dataset = load_dataset_from_zip(zip_file)
+        if dataset_name in datasets:
+            return gr.update(), f"Dataset '{dataset_name}' already exists."
+        datasets[dataset_name] = dataset
+        return gr.Dropdown.update(choices=list(datasets.keys()), value=dataset_name), f"Dataset '{dataset_name}' uploaded."
+
+    upload_button.click(
+        upload_dataset,
+        inputs=[upload_input, datasets],
+        outputs=[dataset_selector, message_box]
+    )
+
+    def select_dataset(dataset_name, datasets):
+        if dataset_name in datasets:
+            dataset = datasets[dataset_name]
+            html_content = display_dataset_html(dataset)
+            return current_dataset_name.update(value=dataset_name), dataset_html.update(value=html_content), ""
+        else:
+            return current_dataset_name.update(value=""), dataset_html.update(value="<div>Select a dataset.</div>"), ""
+
+    dataset_selector.change(
+        select_dataset,
+        inputs=[dataset_selector, datasets],
+        outputs=[current_dataset_name, dataset_html, message_box]
+    )
+
+    with gr.Tab("Add Entry"):
+        with gr.Row():
+            image_input = gr.Image(label="Upload Image")
+            prompt_input = gr.Textbox(label="Prompt")
+        add_button = gr.Button("Add Entry")
+
+    def add_entry(image_data, prompt, current_dataset_name, datasets):
+        if not current_dataset_name:
+            return datasets, dataset_html.update(), "No dataset selected."
+        if image_data is None or not prompt:
+            return datasets, dataset_html.update(), "Please provide both an image and a prompt."
+        datasets[current_dataset_name].append({'image': image_data, 'prompt': prompt})
+        html_content = display_dataset_html(datasets[current_dataset_name])
+        return datasets, dataset_html.update(value=html_content), f"Entry added to dataset '{current_dataset_name}'."
+
+    add_button.click(
+        add_entry,
+        inputs=[image_input, prompt_input, current_dataset_name, datasets],
+        outputs=[datasets, dataset_html, message_box]
+    )
+
+    with gr.Tab("Edit / Delete Entry"):
+        index_input = gr.Number(label="Entry Index", precision=0)
+        new_prompt_input = gr.Textbox(label="New Prompt (for Edit)")
+        with gr.Row():
+            edit_button = gr.Button("Edit Entry")
+            delete_button = gr.Button("Delete Entry")
+
+    def edit_entry(index, new_prompt, current_dataset_name, datasets):
+        if not current_dataset_name:
+            return datasets, dataset_html.update(), "No dataset selected."
+        if index is None or not new_prompt:
+            return datasets, dataset_html.update(), "Please provide both index and new prompt."
+        index = int(index)
+        if 0 <= index < len(datasets[current_dataset_name]):
+            datasets[current_dataset_name][index]['prompt'] = new_prompt
+            html_content = display_dataset_html(datasets[current_dataset_name])
+            return datasets, dataset_html.update(value=html_content), f"Entry {index} updated."
+        else:
+            return datasets, dataset_html.update(), "Invalid index."
+
+    edit_button.click(
+        edit_entry,
+        inputs=[index_input, new_prompt_input, current_dataset_name, datasets],
+        outputs=[datasets, dataset_html, message_box]
+    )
+
+    def delete_entry(index, current_dataset_name, datasets):
+        if not current_dataset_name:
+            return datasets, dataset_html.update(), "No dataset selected."
+        if index is None:
+            return datasets, dataset_html.update(), "Please provide an index."
+        index = int(index)
+        if 0 <= index < len(datasets[current_dataset_name]):
+            del datasets[current_dataset_name][index]
+            html_content = display_dataset_html(datasets[current_dataset_name])
+            return datasets, dataset_html.update(value=html_content), f"Entry {index} deleted."
+        else:
+            return datasets, dataset_html.update(), "Invalid index."
+
+    delete_button.click(
+        delete_entry,
+        inputs=[index_input, current_dataset_name, datasets],
+        outputs=[datasets, dataset_html, message_box]
+    )
+
+    with gr.Tab("Download Dataset"):
+        download_button = gr.Button("Download Dataset")
+        download_output = gr.File(label="Download Zip")
+        def download_dataset(current_dataset_name, datasets):
+            if not current_dataset_name:
+                return None, "No dataset selected."
+            zip_buffer = save_dataset_to_zip(current_dataset_name, datasets[current_dataset_name])
+            return zip_buffer.getvalue(), f"Dataset '{current_dataset_name}' is ready for download."
+        download_button.click(
+            download_dataset,
+            inputs=[current_dataset_name, datasets],
+            outputs=[download_output, message_box]
+        )
+
+    # Initially update dataset_html
+    demo.load(
+        lambda datasets: gr.Dropdown.update(choices=list(datasets.keys())),
+        inputs=None,
+        outputs=dataset_selector
+    )
+
+demo.launch()