Update app.py

app.py

@@ -76,7 +76,6 @@ def load_dataset_from_zip(zip_file_path):
                 dataset_name = dataset_name_guess
                 dataset_path = temp_dir
 
-        images_dir = os.path.join(dataset_path, 'images')
         annotations_path = os.path.join(dataset_path, 'annotations.jsonl')
         dataset = []
 
@@ -113,7 +112,6 @@ def display_dataset_html(dataset, page_number=0, items_per_page=5):
         start_idx = page_number * items_per_page
         end_idx = start_idx + items_per_page
         dataset_slice = dataset[start_idx:end_idx]
-        total_pages = (len(dataset) - 1) // items_per_page + 1
         html_content = '''
         <div style="display: flex; overflow-x: auto; padding: 10px; border: 1px solid #ccc;">
         '''
@@ -142,47 +140,179 @@ with gr.Blocks() as demo:
     entry_selector = gr.Dropdown(label="Select Entry to Edit/Delete")
     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")
+    # Tabs as before
+    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]
+            )
+
+        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 add_entry(image_data, prompt, current_dataset_name, datasets):
+                if not current_dataset_name:
+                    return datasets, gr.update(), gr.update(), gr.update(), "No dataset selected."
+                if image_data is None or not prompt:
+                    return datasets, gr.update(), 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 when a new entry is added
+                page_number = 0
+                html_content = display_dataset_html(dataset, page_number=page_number)
+                # Update entry_selector options
+                entry_options = [f"{idx}: {entry['prompt'][:30]}" for idx, entry in enumerate(dataset)]
+                return datasets, page_number, gr.update(value=html_content), gr.update(choices=entry_options), f"Entry added to dataset '{current_dataset_name}'."
+
+            add_button.click(
+                add_entry,
+                inputs=[image_input, prompt_input, current_dataset_name, datasets],
+                outputs=[datasets, current_page_number, gr.HTML(), entry_selector, message_box]
+            )
+
+        with gr.TabItem("Edit / Delete Entry"):
             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]
-    )
+                selected_image = gr.Image(label="Selected Image", interactive=False, type="numpy")
+                selected_prompt = gr.Textbox(label="Current Prompt", interactive=False)
+                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 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_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(), gr.update(), "No dataset selected."
+                if not entry_option or not new_prompt.strip():
+                    return datasets, gr.update(), gr.update(), gr.update(), gr.update(), "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=""), gr.update(), f"Entry {index} updated."
+
+            edit_button.click(
+                edit_entry,
+                inputs=[entry_selector, new_prompt_input, current_dataset_name, datasets, current_page_number],
+                outputs=[datasets, gr.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(), gr.update(), "No dataset selected."
+                if not entry_option:
+                    return datasets, gr.update(), 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), gr.update(value=""), message_box, f"Entry {index} deleted."
+
+            delete_button.click(
+                delete_entry,
+                inputs=[entry_selector, current_dataset_name, datasets, current_page_number],
+                outputs=[datasets, gr.HTML(), entry_selector, selected_image, selected_prompt, message_box]
+            )
+
+        with gr.TabItem("Download Dataset"):
+            download_button = gr.Button("Download Dataset")
+            download_output = gr.File(label="Download Zip", interactive=False)
+
+            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."
+
+            download_button.click(
+                download_dataset,
+                inputs=[current_dataset_name, datasets],
+                outputs=[download_output, 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]
-    )
+    # 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")
 
     def select_dataset(dataset_name, datasets):
         if dataset_name in datasets:
@@ -197,137 +327,9 @@ with gr.Blocks() as demo:
     dataset_selector.change(
         select_dataset,
         inputs=[dataset_selector, datasets],
-        outputs=[current_dataset_name, current_page_number, gr.HTML(), entry_selector, message_box]
+        outputs=[current_dataset_name, current_page_number, dataset_html, entry_selector, message_box]
     )
 
-    with gr.Tab("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 add_entry(image_data, prompt, current_dataset_name, datasets):
-        if not current_dataset_name:
-            return datasets, gr.update(), gr.update(), gr.update(), "No dataset selected."
-        if image_data is None or not prompt:
-            return datasets, gr.update(), 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 when a new entry is added
-        page_number = 0
-        html_content = display_dataset_html(dataset, page_number=page_number)
-        # Update entry_selector options
-        entry_options = [f"{idx}: {entry['prompt'][:30]}" for idx, entry in enumerate(dataset)]
-        return datasets, page_number, gr.update(value=html_content), gr.update(choices=entry_options), f"Entry added to dataset '{current_dataset_name}'."
-
-    add_button.click(
-        add_entry,
-        inputs=[image_input, prompt_input, current_dataset_name, datasets],
-        outputs=[datasets, current_page_number, gr.HTML(), entry_selector, message_box]
-    )
-
-    with gr.Tab("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)
-            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 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_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(), gr.update(), "No dataset selected."
-        if not entry_option or not new_prompt.strip():
-            return datasets, gr.update(), gr.update(), gr.update(), gr.update(), "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=""), gr.update(), f"Entry {index} updated."
-
-    edit_button.click(
-        edit_entry,
-        inputs=[entry_selector, new_prompt_input, current_dataset_name, datasets, current_page_number],
-        outputs=[datasets, gr.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(), gr.update(), "No dataset selected."
-        if not entry_option:
-            return datasets, gr.update(), 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), gr.update(value=""), message_box, f"Entry {index} deleted."
-
-    delete_button.click(
-        delete_entry,
-        inputs=[entry_selector, current_dataset_name, datasets, current_page_number],
-        outputs=[datasets, gr.HTML(), entry_selector, selected_image, selected_prompt, message_box]
-    )
-
-    with gr.Tab("Download Dataset"):
-        download_button = gr.Button("Download Dataset")
-        download_output = gr.File(label="Download Zip", interactive=False)
-        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."
-
-        download_button.click(
-        download_dataset,
-        inputs=[current_dataset_name, datasets],
-        outputs=[download_output, message_box]
-    )
-
-    # 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")
-
     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."