Update app.py

app.py

@@ -1,393 +1,91 @@
+># Import necessary libraries
 import gradio as gr
+import json
 import os
 import zipfile
-import json
-from io import BytesIO
-import base64
-from PIL import Image
-import uuid
-import tempfile
-import numpy as np
-
-def save_dataset_to_zip(dataset_name, dataset):
-    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)
-        # Decode the base64 image data
-        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 with the dataset_name as the top-level folder
-    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, temp_dir)
-                zipf.write(abs_file, rel_file)
-
-    zip_buffer.seek(0)
-    return zip_buffer
-
-def load_dataset_from_zip(zip_file_path):
-    temp_dir = tempfile.mkdtemp()
-    try:
-        with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
-            zip_ref.extractall(temp_dir)
-
-        # Get dataset name from zip file name
-        dataset_name_guess = os.path.splitext(os.path.basename(zip_file_path))[0]
-        dataset_path = os.path.join(temp_dir, dataset_name_guess)
-
-        if os.path.exists(dataset_path):
-            dataset_name = dataset_name_guess
-        else:
-            # If the dataset_name directory doesn't exist, try to find the top-level directory
-            entries = [entry for entry in os.listdir(temp_dir) if os.path.isdir(os.path.join(temp_dir, entry))]
-            if entries:
-                dataset_name = entries[0]
-                dataset_path = os.path.join(temp_dir, dataset_name)
-            else:
-                # Files are directly in temp_dir
-                dataset_name = dataset_name_guess
-                dataset_path = temp_dir
-
-        annotations_path = os.path.join(dataset_path, 'annotations.jsonl')
-        dataset = []
-
-        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
-                    })
-        else:
-            # If annotations file not found
-            return None, []
-
-        return dataset_name, dataset
-    except Exception as e:
-        print(f"Error loading dataset: {e}")
-        return None, []
-
-def display_dataset_html(dataset, page_number=0, items_per_page=2):
-    if dataset:
-        start_idx = page_number * items_per_page
-        end_idx = start_idx + items_per_page
-        dataset_slice = dataset[start_idx:end_idx]
-        html_content = '''
-        <div style="display: flex; overflow-x: auto; padding: 10px; border: 1px solid #ccc;">
-        '''
-        for idx_offset, entry in enumerate(dataset_slice):
-            idx = start_idx + idx_offset
-            image_data = entry['image']
-            prompt = entry['prompt']
-            html_content += f"""
-            <div style="display: flex; flex-direction: column; align-items: center; margin-right: 20px;">
-                <div style="margin-bottom: 5px;">{idx}</div>
-                <img src="{image_data}" alt="Image {idx}" style="max-height: 150px;"/>
-                <div style="max-width: 150px; word-wrap: break-word; text-align: center;">{prompt}</div>
-            </div>
-            """
-        html_content += '</div>'
-        return html_content
-    else:
-        return "<div>No entries in dataset.</div>"
-
-#Interface
-with gr.Blocks() as demo:
-    gr.Markdown("<h1 style='text-align: center; margin-bottom: 1px;'>Dataset Creator</h1>")
-    gr.Markdown("You must create/upload a dataset before selecting one")
-    datasets = gr.State({})
-    current_dataset_name = gr.State("")
-    current_page_number = gr.State(0)
-
-    # Top-level components
-    with gr.Column():
-        dataset_selector = gr.Dropdown(label="Select Dataset", interactive=True)
-        message_box = gr.Textbox(interactive=False, label="Message")
-
-    # Dataset Viewer and Pagination Controls at the Bottom
-    with gr.Column():
-        gr.Markdown("### Dataset Viewer")
-        dataset_html = gr.HTML()
-        with gr.Row():
-            prev_button = gr.Button("Previous Page")
-            next_button = gr.Button("Next Page")
-
-    # Tabs
-    with gr.Tabs():
-        with gr.TabItem("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", type="filepath", 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.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_path, datasets):
-                if not zip_file_path:
-                    return gr.update(), "Please upload a zip file."
-                dataset_name, dataset = load_dataset_from_zip(zip_file_path)
-                if dataset_name is None:
-                    return gr.update(), "Failed to load dataset from zip file."
-                if dataset_name in datasets:
-                    return gr.update(), f"Dataset '{dataset_name}' already exists."
-                datasets[dataset_name] = dataset
-                return gr.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]
-            )
+# Define helper functions
 
-        with gr.TabItem("Add Entry"):
-            with gr.Row():
-                image_input = gr.Image(label="Upload Image", type="numpy")
-                prompt_input = gr.Textbox(label="Prompt")
-            add_button = gr.Button("Add Entry")
+def create_dataset(dataset_name):
+    dataset_path = f'{dataset_name}.zip'
+    if not os.path.exists(dataset_path):
+        with zipfile.ZipFile(dataset_path, 'w') as zip_file:
+            zip_file.writestr('images/', '')
+            zip_file.writestr('data.jsonl', '')
 
-            def add_entry(image_data, prompt, current_dataset_name, datasets):
-                if not current_dataset_name:
-                    return datasets, gr.update(), gr.update(), "No dataset selected."
-                if image_data is None or not prompt:
-                    return datasets, gr.update(), gr.update(), "Please provide both an image and a prompt."
-                # Convert image_data to base64
-                image = Image.fromarray(image_data.astype('uint8'))
-                buffered = BytesIO()
-                image.save(buffered, format="PNG")
-                img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
-                img_data = f"data:image/png;base64,{img_str}"
-                datasets[current_dataset_name].append({'image': img_data, 'prompt': prompt})
-                dataset = datasets[current_dataset_name]
-                # Reset page number to 0 and refresh HTML
-                page_number = 0
-                dataset = datasets[current_dataset_name]
-                html_content = display_dataset_html(dataset, page_number=page_number)
-                return datasets, page_number, gr.update(value=html_content), f"Entry added to dataset '{current_dataset_name}'."
+    return dataset_path
 
-            add_button.click(
-                add_entry,
-                inputs=[image_input, prompt_input, current_dataset_name, datasets],
-                outputs=[datasets, current_page_number, dataset_html, message_box]
-            )
 
-        with gr.TabItem("Edit / Delete Entry"):
-            with gr.Column():
-                selected_image = gr.Image(label="Selected Image", interactive=False, type="numpy")
-                selected_prompt = gr.Textbox(label="Current Prompt", interactive=False)
-                # Define entry_selector here
-                entry_selector = gr.Dropdown(label="Select Entry to Edit/Delete")
-                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 upload_pair(dataset_path, image, prompt):
+    with zipfile.ZipFile(dataset_path, 'a') as zip_file:
+        image_path = f'images/{image.name}'
+        zip_file.writestr(image_path, image.read())
+        data = {'image': image_path, 'prompt': prompt}
+        zip_file.writestr('data.jsonl', json.dumps(data) + '\n')
 
-            def update_selected_entry(entry_option, current_dataset_name, datasets):
-                if not current_dataset_name or not entry_option:
-                    return gr.update(), gr.update()
-                index = int(entry_option.split(":")[0])
-                entry = datasets[current_dataset_name][index]
-                image_data = entry['image']
-                prompt = entry['prompt']
-                # Decode base64 image data to numpy array
-                image_bytes = base64.b64decode(image_data.split(",")[1])
-                image = Image.open(BytesIO(image_bytes))
-                image_array = np.array(image)
-                return gr.update(value=image_array), gr.update(value=prompt)
 
-            entry_selector.change(
-                update_selected_entry,
-                inputs=[entry_selector, current_dataset_name, datasets],
-                outputs=[selected_image, selected_prompt]
-            )
+def edit_prompt(dataset_path, image_path, new_prompt):
+    with zipfile.ZipFile(dataset_path, 'r') as zip_file:
+        data = json.load(zip_file.open('data.jsonl'))
+        for item in data:
+            if item['image'] == image_path:
+                item['prompt'] = new_prompt
+                break
 
-            def edit_entry(entry_option, new_prompt, current_dataset_name, datasets, current_page_number):
-                if not current_dataset_name:
-                    return datasets, gr.update(), gr.update(), gr.update(), f"No dataset selected."
-                if not entry_option or not new_prompt.strip():
-                    return datasets, gr.update(), gr.update(), gr.update(), f"Please select an entry and provide a new prompt."
-                index = int(entry_option.split(":")[0])
-                datasets[current_dataset_name][index]['prompt'] = new_prompt
-                dataset = datasets[current_dataset_name]
-                html_content = display_dataset_html(dataset, page_number=current_page_number)
-                # Update entry_selector options
-                entry_options = [f"{idx}: {entry['prompt'][:30]}" for idx, entry in enumerate(dataset)]
-                return datasets, gr.update(value=html_content), gr.update(choices=entry_options), gr.update(value=""), f"Entry {index} updated."
+    with zipfile.ZipFile(dataset_path, 'w') as zip_file:
+        zip_file.writestr('data.jsonl', json.dumps(data))
 
-            edit_button.click(
-                edit_entry,
-                inputs=[entry_selector, new_prompt_input, current_dataset_name, datasets, current_page_number],
-                outputs=[datasets, dataset_html, entry_selector, new_prompt_input, message_box]
-            )
 
-            def delete_entry(entry_option, current_dataset_name, datasets, current_page_number):
-                if not current_dataset_name:
-                    return datasets, gr.update(), gr.update(), gr.update(), gr.update(), "No dataset selected."
-                if not entry_option:
-                    return datasets, gr.update(), gr.update(), gr.update(), gr.update(), "Please select an entry to delete."
-                index = int(entry_option.split(":")[0])
-                del datasets[current_dataset_name][index]
-                dataset = datasets[current_dataset_name]
-                html_content = display_dataset_html(dataset, page_number=current_page_number)
-                # Update entry_selector options
-                entry_options = [f"{idx}: {entry['prompt'][:30]}" for idx, entry in enumerate(dataset)]
-                return datasets, gr.update(value=html_content), gr.update(choices=entry_options), gr.update(value=None), f"Entry {index} deleted."
+def delete_pair(dataset_path, image_path):
+    with zipfile.ZipFile(dataset_path, 'r') as zip_file:
+        data = json.load(zip_file.open('data.jsonl'))
+        data = [item for item in data if item['image'] != image_path]
 
-            delete_button.click(
-                delete_entry,
-                inputs=[entry_selector, current_dataset_name, datasets, current_page_number],
-                outputs=[datasets, dataset_html, entry_selector, selected_image, message_box]
-            )
+    with zipfile.ZipFile(dataset_path, 'w') as zip_file:
+        zip_file.writestr('data.jsonl', json.dumps(data))
 
-            # Function to update entry_selector options
-            def update_entry_selector(current_dataset_name, datasets):
-                if current_dataset_name in datasets:
-                    dataset = datasets[current_dataset_name]
-                    entry_options = [f"{idx}: {entry['prompt'][:30]}" for idx, entry in enumerate(dataset)]
-                    return gr.update(choices=entry_options)
-                else:
-                    return gr.update(choices=[])
 
-            # Update entry_selector when dataset is selected
-            dataset_selector.change(
-                update_entry_selector,
-                inputs=[current_dataset_name, datasets],
-                outputs=[entry_selector]
-            )
+def download_dataset(dataset_path):
+    return dataset_path
 
-            # Also update entry_selector when an entry is added in "Add Entry" tab
-            add_button.click(
-                update_entry_selector,
-                inputs=[current_dataset_name, datasets],
-                outputs=[entry_selector]
-            )
+# Define Gradio application
 
-        with gr.TabItem("Download Dataset"):
-            download_button = gr.Button("Download Dataset")
-            download_output = gr.File(label="Download Zip", interactive=False)
+demo = gr.Blocks()
 
-            def download_dataset(current_dataset_name, datasets):
-                if not current_dataset_name:
-                    return None, "No dataset selected."
-                if not datasets[current_dataset_name]:
-                    return None, "Dataset is empty."
-                zip_buffer = save_dataset_to_zip(current_dataset_name, datasets[current_dataset_name])
-                # Write zip_buffer to a temporary file
-                temp_dir = tempfile.mkdtemp()
-                zip_path = os.path.join(temp_dir, f"{current_dataset_name}.zip")
-                with open(zip_path, 'wb') as f:
-                    f.write(zip_buffer.getvalue())
-                return zip_path, f"Dataset '{current_dataset_name}' is ready for download."
+with demo:
+    # Create dataset
+    dataset_name = gr.Textbox(label='Dataset Name')
+    create_button = gr.Button('Create Dataset')
+    create_button.click(create_dataset, inputs=[dataset_name], outputs=[])
 
-            download_button.click(
-                download_dataset,
-                inputs=[current_dataset_name, datasets],
-                outputs=[download_output, message_box]
-            )
+    # Upload pair
+    image_upload = gr.File(label='Image')
+    prompt = gr.Textbox(label='Prompt')
+    upload_button = gr.Button('Upload Pair')
+    upload_button.click(upload_pair, inputs=[dataset_name, image_upload, prompt], outputs=[])
 
-    def select_dataset(dataset_name, datasets):
-        if dataset_name in datasets:
-            dataset = datasets[dataset_name]
-            html_content = display_dataset_html(dataset, page_number=0)
-            return dataset_name, 0, gr.update(value=html_content), f"Dataset '{dataset_name}' selected."
-        else:
-            return "", 0, gr.update(value="<div>Select a dataset.</div>"), ""
+    # Edit prompt
+    image_path = gr.Textbox(label='Image Path')
+    new_prompt = gr.Textbox(label='New Prompt')
+    edit_button = gr.Button('Edit Prompt')
+    edit_button.click(edit_prompt, inputs=[dataset_name, image_path, new_prompt], outputs=[])
 
-    dataset_selector.change(
-        select_dataset,
-        inputs=[dataset_selector, datasets],
-        outputs=[current_dataset_name, current_page_number, dataset_html, message_box]
-    )
-            
-    def change_page(action, current_page_number, datasets, current_dataset_name):
-        if not current_dataset_name:
-            return current_page_number, gr.update(), "No dataset selected."
-        dataset = datasets[current_dataset_name]
-        total_pages = (len(dataset) - 1) // 5 + 1
-        if action == "next":
-            if current_page_number + 1 < total_pages:
-                current_page_number += 1
-        elif action == "prev":
-            if current_page_number > 0:
-                current_page_number -= 1
-        html_content = display_dataset_html(dataset, page_number=current_page_number)
-        return current_page_number, gr.update(value=html_content), ""
-    
-    prev_button.click(
-        fn=lambda current_page_number, datasets, current_dataset_name: change_page("prev", current_page_number, datasets, current_dataset_name),
-        inputs=[current_page_number, datasets, current_dataset_name],
-        outputs=[current_page_number, dataset_html, message_box]
-    )
+    # Delete pair
+    delete_button = gr.Button('Delete Pair')
+    delete_button.click(delete_pair, inputs=[dataset_name, image_path], outputs=[])
 
-    next_button.click(
-        fn=lambda current_page_number, datasets, current_dataset_name: change_page("next", current_page_number, datasets, current_dataset_name),
-        inputs=[current_page_number, datasets, current_dataset_name],
-        outputs=[current_page_number, dataset_html, message_box]
-    )
+    # Download dataset
+    download_button = gr.Button('Download Dataset')
+    download_button.click(download_dataset, inputs=[dataset_name], outputs=[])
 
-    # Initialize dataset_selector
-    def initialize_components(datasets):
-        return gr.update(choices=list(datasets.keys()))
+    # Upload dataset
+    dataset_upload = gr.File(label='Dataset')
+    upload_dataset_button = gr.Button('Upload Dataset')
+    upload_dataset_button.click(create_dataset, inputs=[dataset_upload], outputs=[])
 
-    demo.load(
-        initialize_components,
-        inputs=[datasets],
-        outputs=[dataset_selector]
-    )
+    # Horizontal gallery
+    gallery = gr.Gallery(label='Dataset Gallery')
+    demo.append(gallery)
 
+# Launch Gradio application
 demo.launch()