Update app.py

# Key Changes and Explanations:

## get_images(): This function now takes a batch_size argument and fetches multiple images and their metadata, similar to the Dataset Viewer. It uses the get_next_sample() function to get images one by one until the batch size is reached or the dataset is exhausted.

## save_labeled_data(): This function is updated to work similarly to get_images(), but it processes only one image at a time. It gets the next sample using get_next_sample(), saves the labeled data, and then updates the UI with the new image and metadata.

## labeling_ui():

Uses gr.Gallery to display images, aligning with the Dataset Viewer's approach.

Uses gr.DataFrame for metadata, which is better suited for structured data.

The initialize_labeling_ui() function is similar to before but uses the new get_images() to fetch the first image and its metadata.

Button click events are updated to work with the new image retrieval logic.

## display_ui(): Remains mostly the same, as it was already functional.

## State Management: Uses gr.State to store the current sample, which is updated each time a new sample is fetched for labeling.

app.py

@@ -5,6 +5,8 @@ import numpy as np
 import random
 import os
 import json
+import utils
+from pandas import DataFrame
 
 # --- Configuration ---
 DATASET_SUBSET = "satellogic"
@@ -12,6 +14,7 @@ NUM_SAMPLES_TO_LABEL = 100  # You can adjust this
 LABELED_DATA_FILE = "labeled_data.json"
 DISPLAY_N_COOL = 5  # How many "cool" examples to display alongside each new image.
 SAMPLE_SEED = 10  # The seed to use when sampling the dataset for the demo.
+BATCH_SIZE = 10
 
 # --- Load Dataset ---
 dataset = ev.load_dataset(DATASET_SUBSET, shards=[SAMPLE_SEED])
@@ -33,64 +36,104 @@ def get_next_sample():
     try:
         sample = next(data_iter)
         sample = ev.item_to_images(DATASET_SUBSET, sample)
-        image = sample["rgb"][0]
-        metadata = sample["metadata"]
-        return image, metadata, len(labeled_data)
+        return sample
     except StopIteration:
         print("No more samples in the dataset.")
-        return None, None, len(labeled_data)
+        return None
+
+def get_images(batch_size, state):
+    subset = state["subset"]
+    images = []
+    metadatas = []
+    for _ in range(batch_size):
+      sample = get_next_sample()
+      if sample is None:
+        break
+
+      image = sample["rgb"][0]
+      metadata = sample["metadata"]
+      metadata["map"] = f'<a href="{utils.get_google_map_link(sample, subset)}" target="_blank">🧭</a>'
+
+      images.append(image)
+      metadatas.append(metadata)
+
+    return images, DataFrame(metadatas)
 
 # --- Save Labeled Data ---
-def save_labeled_data(image, metadata, label):
+def save_labeled_data(label, state):
     global labeled_data
-    image_bytes = image.convert("RGB").tobytes() if image else None
+    sample = state["sample"]
+    if sample is None:
+        return "No image to label", None, DataFrame()
+
+    image = sample["rgb"][0]  # Get the PIL Image object
+
+    image_bytes = image.convert("RGB").tobytes()
     labeled_data.append({
         "image": image_bytes,
-        "metadata": metadata,
+        "metadata": sample["metadata"],
         "label": label
     })
+
     with open(LABELED_DATA_FILE, "w") as f:
         json.dump(labeled_data, f)
-    new_image, new_metadata, count = get_next_sample()
-    if new_image is None:
-        return "Dataset exhausted.", None, "", f"Labeled {count} samples."
-    return "", new_image, json.dumps(new_metadata["bounds"]), f"Labeled {count} samples."
+
+    new_sample = get_next_sample()
+    if new_sample is None:
+      state["sample"] = None
+      return "Dataset exhausted.", None, DataFrame()
+
+    state["sample"] = new_sample
+
+    new_image = new_sample["rgb"][0]
+    new_metadata = new_sample["metadata"]
+    new_metadata["map"] = f'<a href="{utils.get_google_map_link(new_sample, DATASET_SUBSET)}" target="_blank">🧭</a>'
+    
+    return "", new_image, DataFrame([new_metadata])
 
 # --- Gradio Interface ---
 # --- Labeling UI ---
 def labeling_ui():
+    
+    state = gr.State({"sample": None, "subset": DATASET_SUBSET})
+
     with gr.Row():
         with gr.Column():
-            image_component = gr.Image(label="Satellite Image", type="pil")
-            metadata_text = gr.Textbox(label="Metadata (Bounds)")
-            label_count_text = gr.Textbox(label="Label Count")
+            gallery = gr.Gallery(label="Satellite Image", interactive=False, columns=1, object_fit="scale-down")
+
             with gr.Row():
                 cool_button = gr.Button("Cool")
                 not_cool_button = gr.Button("Not Cool")
+
+            table = gr.DataFrame(datatype="html")
+
+            def initialize_labeling_ui():
+              sample = get_next_sample()
+              
+              image, metadata = get_images(1, {"sample": None, "subset": DATASET_SUBSET})
+
+              return sample, image, metadata
+
+            initial_sample, initial_image, initial_metadata = initialize_labeling_ui()
+            gallery.value = initial_image
+            table.value = initial_metadata
+            state.value["sample"] = initial_sample
+
+            # Handle button clicks
             cool_button.click(
-                fn=save_labeled_data,
-                inputs=[image_component, metadata_text],
-                outputs=[gr.Textbox(label="Debug"), image_component, metadata_text, label_count_text]
+                fn=lambda label, state: save_labeled_data(label, state),
+                inputs=[gr.Textbox(visible=False, value="cool"), state],
+                outputs=[gr.Textbox(label="Debug"), gallery, table]
             )
             not_cool_button.click(
-                fn=save_labeled_data,
-                inputs=[image_component, metadata_text],
-                outputs=[gr.Textbox(label="Debug"), image_component, metadata_text, label_count_text]
+                fn=lambda label, state: save_labeled_data(label, state),
+                inputs=[gr.Textbox(visible=False, value="not cool"), state],
+                outputs=[gr.Textbox(label="Debug"), gallery, table]
             )
 
-    def initialize_labeling_ui():
-        image, metadata, count = get_next_sample()
-        if image:
-            return image, json.dumps(metadata["bounds"]), f"Labeled {count} samples."
-        return None, "", "No samples loaded."
-
-    initial_image, initial_metadata, initial_count = initialize_labeling_ui()
-    image_component.value = initial_image
-    metadata_text.value = initial_metadata
-    label_count_text.value = initial_count
-
 # --- Display UI ---
 def display_ui():
+
     def get_random_cool_images(n):
         cool_samples = [d for d in labeled_data if d["label"] == "cool"]
         return [Image.frombytes("RGB", (384, 384), s["image"]) for s in cool_samples] if len(cool_samples) >= n else []