Update app.py

app.py

@@ -108,19 +108,27 @@ def load_dataset_from_zip(zip_file_path):
         print(f"Error loading dataset: {e}")
         return None, []
 
-def display_dataset_html(dataset):
+def display_dataset_html(dataset, page_number=0, items_per_page=5):
     if dataset:
-        html_content = ""
-        for idx, entry in enumerate(dataset):
+        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;">
+        '''
+        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; 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 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>"
@@ -129,9 +137,9 @@ with gr.Blocks() as demo:
     gr.Markdown("<h1 style='text-align: center; margin-bottom: 20px;'>Dataset Builder</h1>")
     datasets = gr.State({})
     current_dataset_name = gr.State("")
+    current_page_number = gr.State(0)
     dataset_selector = gr.Dropdown(label="Select Dataset", interactive=True)
     entry_selector = gr.Dropdown(label="Select Entry to Edit/Delete")
-    dataset_html = gr.HTML()
     message_box = gr.Textbox(interactive=False, label="Message")
 
     with gr.Tab("Create / Upload Dataset"):
@@ -179,17 +187,17 @@ with gr.Blocks() as demo:
     def select_dataset(dataset_name, datasets):
         if dataset_name in datasets:
             dataset = datasets[dataset_name]
-            html_content = display_dataset_html(dataset)
+            html_content = display_dataset_html(dataset, page_number=0)
             # Update entry_selector options
             entry_options = [f"{idx}: {entry['prompt'][:30]}" for idx, entry in enumerate(dataset)]
-            return dataset_name, gr.update(value=html_content), gr.update(choices=entry_options), ""
+            return dataset_name, 0, gr.update(value=html_content), gr.update(choices=entry_options), ""
         else:
-            return "", gr.update(value="<div>Select a dataset.</div>"), gr.update(choices=[]), ""
+            return "", 0, gr.update(value="<div>Select a dataset.</div>"), gr.update(choices=[]), ""
 
     dataset_selector.change(
         select_dataset,
         inputs=[dataset_selector, datasets],
-        outputs=[current_dataset_name, dataset_html, entry_selector, message_box]
+        outputs=[current_dataset_name, current_page_number, gr.HTML(), entry_selector, message_box]
     )
 
     with gr.Tab("Add Entry"):
@@ -200,9 +208,9 @@ with gr.Blocks() as demo:
 
     def add_entry(image_data, prompt, current_dataset_name, datasets):
         if not current_dataset_name:
-            return datasets, gr.update(), gr.update(), "No dataset selected."
+            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(), "Please provide both an image and a 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()
@@ -211,15 +219,17 @@ with gr.Blocks() as demo:
         img_data = f"data:image/png;base64,{img_str}"
         datasets[current_dataset_name].append({'image': img_data, 'prompt': prompt})
         dataset = datasets[current_dataset_name]
-        html_content = display_dataset_html(dataset)
+        # 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, gr.update(value=html_content), gr.update(choices=entry_options), f"Entry added to dataset '{current_dataset_name}'."
+        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, dataset_html, entry_selector, message_box]
+        outputs=[datasets, current_page_number, gr.HTML(), entry_selector, message_box]
     )
 
     with gr.Tab("Edit / Delete Entry"):
@@ -250,42 +260,42 @@ with gr.Blocks() as demo:
         outputs=[selected_image, selected_prompt]
     )
 
-    def edit_entry(entry_option, new_prompt, current_dataset_name, datasets):
+    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(), "No dataset selected."
+            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(), "Please select an entry and provide a new prompt."
+            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)
+        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."
+        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],
-        outputs=[datasets, dataset_html, entry_selector, new_prompt_input, message_box]
+        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):
+    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."
+            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(), "Please select an entry to delete."
+            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)
+        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=""), f"Entry {index} deleted."
+        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],
-        outputs=[datasets, dataset_html, entry_selector, selected_image, selected_prompt, message_box]
+        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"):
@@ -310,6 +320,40 @@ with gr.Blocks() as demo:
             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."
+        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(
+        change_page,
+        inputs=["prev", current_page_number, datasets, current_dataset_name],
+        outputs=[current_page_number, dataset_html, message_box]
+    )
+
+    next_button.click(
+        change_page,
+        inputs=["next", current_page_number, datasets, current_dataset_name],
+        outputs=[current_page_number, dataset_html, message_box]
+    )
+
     # Initialize dataset_selector and entry_selector
     def initialize_components(datasets):
         return gr.update(choices=list(datasets.keys())), gr.update(choices=[])