Update app.py

app.py

@@ -7,6 +7,7 @@ import base64
 from PIL import Image
 import uuid
 import tempfile
+import numpy as np
 
 def save_dataset_to_zip(dataset_name, dataset):
     # Create a temporary directory
@@ -22,6 +23,7 @@ def save_dataset_to_zip(dataset_name, dataset):
         # 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
@@ -49,7 +51,6 @@ 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 = []
@@ -143,14 +144,16 @@ with gr.Blocks() as demo:
         if dataset_name in datasets:
             dataset = datasets[dataset_name]
             html_content = display_dataset_html(dataset)
-            return dataset_name, gr.update(value=html_content), ""
+            # 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), ""
         else:
-            return "", gr.update(value="<div>Select a dataset.</div>"), ""
+            return "", 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, message_box]
+        outputs=[current_dataset_name, dataset_html, gr.Variable(), message_box]
     )
 
     with gr.Tab("Add Entry"):
@@ -161,62 +164,89 @@ with gr.Blocks() as demo:
 
     def add_entry(image_data, prompt, current_dataset_name, datasets):
         if not current_dataset_name:
-            return datasets, gr.update(), "No dataset selected."
+            return datasets, gr.update(), gr.update(), "No dataset selected."
         if image_data is None or not prompt:
-            return datasets, gr.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, gr.update(value=html_content), f"Entry added to dataset '{current_dataset_name}'."
+            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]
+        html_content = display_dataset_html(dataset)
+        # 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}'."
 
     add_button.click(
         add_entry,
         inputs=[image_input, prompt_input, current_dataset_name, datasets],
-        outputs=[datasets, dataset_html, message_box]
+        outputs=[datasets, dataset_html, gr.Variable(), message_box]
     )
 
     with gr.Tab("Edit / Delete Entry"):
-        index_input = gr.Number(label="Entry Index", precision=0)
+        # Entry selector
+        entry_selector = gr.Dropdown(label="Select Entry to Edit/Delete")
+        selected_image = gr.Image(label="Selected Image", interactive=False)
+        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 edit_entry(index, new_prompt, current_dataset_name, datasets):
+    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']
+        return gr.update(value=image_data), 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):
         if not current_dataset_name:
-            return datasets, gr.update(), "No dataset selected."
-        if index is None or new_prompt is None or new_prompt.strip() == '':
-            return datasets, gr.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, gr.update(value=html_content), f"Entry {index} updated."
-        else:
-            return datasets, gr.update(), "Invalid index."
+            return datasets, gr.update(), gr.update(), "No dataset selected."
+        if not entry_option or not new_prompt.strip():
+            return datasets, 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)
+        # 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 {index} updated."
 
     edit_button.click(
         edit_entry,
-        inputs=[index_input, new_prompt_input, current_dataset_name, datasets],
-        outputs=[datasets, dataset_html, message_box]
+        inputs=[entry_selector, new_prompt_input, current_dataset_name, datasets],
+        outputs=[datasets, dataset_html, entry_selector, message_box]
     )
 
-    def delete_entry(index, current_dataset_name, datasets):
+    def delete_entry(entry_option, current_dataset_name, datasets):
         if not current_dataset_name:
-            return datasets, gr.update(), "No dataset selected."
-        if index is None:
-            return datasets, gr.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, gr.update(value=html_content), f"Entry {index} deleted."
-        else:
-            return datasets, gr.update(), "Invalid index."
+            return datasets, gr.update(), gr.update(), "No dataset selected."
+        if not entry_option:
+            return datasets, 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)
+        # 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 {index} deleted."
 
     delete_button.click(
         delete_entry,
-        inputs=[index_input, current_dataset_name, datasets],
-        outputs=[datasets, dataset_html, message_box]
+        inputs=[entry_selector, current_dataset_name, datasets],
+        outputs=[datasets, dataset_html, entry_selector, message_box]
     )
 
     with gr.Tab("Download Dataset"):
@@ -225,6 +255,8 @@ with gr.Blocks() as demo:
         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])
             return zip_buffer.getvalue(), f"Dataset '{current_dataset_name}' is ready for download."
         download_button.click(
@@ -233,14 +265,14 @@ with gr.Blocks() as demo:
             outputs=[download_output, message_box]
         )
 
-    # Initially update dataset_selector
-    def initialize_datasets(datasets):
-        return gr.update(choices=list(datasets.keys()))
-    
+    # Initialize dataset_selector and entry_selector
+    def initialize_components(datasets):
+        return gr.update(choices=list(datasets.keys())), gr.update(choices=[])
+
     demo.load(
-        initialize_datasets,
+        initialize_components,
         inputs=[datasets],
-        outputs=[dataset_selector]
+        outputs=[dataset_selector, entry_selector]
     )
 
 demo.launch()