Update app.py

# Key Changes:

## Image Handling:

The get_next_sample() function now correctly returns a PIL Image object for the image and the raw metadata as extracted from the sample.

The save_labeled_data function takes care of converting the image to bytes (if not None) and stores it along with the metadata.

## Metadata Handling:

Instead of trying to directly assign the metadata to the metadata_text textbox, which might have caused serialization issues if the metadata was not a simple string, we now only pass the "bounds" part of the metadata. This is converted to a JSON string using json.dumps(). This ensures that the textbox always receives a string.

## Initialization of UI Components:

The initialization functions for both labeling_ui and display_ui now properly return the initial values, which are then used to set the .value attribute of the components directly.

## Error Handling for Dataset Exhaustion:

Added a print statement in get_next_sample() and get_new_unlabeled_image() to indicate when the dataset is exhausted. Also, save_labeled_data now returns a specific message when there are no more samples.

## Display UI Initialization:

The get_random_cool_images function in display_ui now handles cases where there are fewer cool samples than requested.

The refresh_display function ensures that cool_images is always a list, even if empty.

# Explanation of Changes:

By explicitly converting the metadata to a JSON string, we avoid passing complex objects to Gradio components that might not be able to handle them.

Returning the PIL Image object directly from functions that interact with Gradio's gr.Image component ensures that the component receives the expected data type.

Initializing the components' values directly using .value during the setup phase ensures that they start with the correct content.

app.py

@@ -12,6 +12,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.
+
 # --- Load Dataset ---
 dataset = ev.load_dataset(DATASET_SUBSET, shards=[SAMPLE_SEED])
 data_iter = iter(dataset)
@@ -34,39 +35,28 @@ def get_next_sample():
         sample = ev.item_to_images(DATASET_SUBSET, sample)
         image = sample["rgb"][0]
         metadata = sample["metadata"]
-
         return image, metadata, len(labeled_data)
     except StopIteration:
-        return None, None, None
+        print("No more samples in the dataset.")
+        return None, None, len(labeled_data)
 
 # --- Save Labeled Data ---
 def save_labeled_data(image, metadata, label):
     global labeled_data
-
-    # Convert PIL Image to bytes before saving
-    if image is not None:
-        image_bytes = image.convert("RGB").tobytes()
-    else:
-        image_bytes = None
-
+    image_bytes = image.convert("RGB").tobytes() if image else None
     labeled_data.append({
         "image": image_bytes,
         "metadata": metadata,
         "label": label
     })
-
     with open(LABELED_DATA_FILE, "w") as f:
         json.dump(labeled_data, f)
-
-    image, metadata, count = get_next_sample()
-
-    if image is None:
-        return "No more samples", gr.Image.update(value=None), "", f"Labeled {count} samples."
-
-    return "", image, str(metadata["bounds"]), f"Labeled {count} samples."
+    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."
 
 # --- Gradio Interface ---
-
 # --- Labeling UI ---
 def labeling_ui():
     with gr.Row():
@@ -74,30 +64,25 @@ def labeling_ui():
             image_component = gr.Image(label="Satellite Image", type="pil")
             metadata_text = gr.Textbox(label="Metadata (Bounds)")
             label_count_text = gr.Textbox(label="Label Count")
-
             with gr.Row():
                 cool_button = gr.Button("Cool")
                 not_cool_button = gr.Button("Not Cool")
-
-            # Handle button clicks
             cool_button.click(
-                fn=lambda image, metadata: save_labeled_data(image, metadata, "cool"),
+                fn=save_labeled_data,
                 inputs=[image_component, metadata_text],
                 outputs=[gr.Textbox(label="Debug"), image_component, metadata_text, label_count_text]
             )
             not_cool_button.click(
-                fn=lambda image, metadata: save_labeled_data(image, metadata, "not cool"),
+                fn=save_labeled_data,
                 inputs=[image_component, metadata_text],
                 outputs=[gr.Textbox(label="Debug"), image_component, metadata_text, label_count_text]
             )
 
-    # Initialize with the first sample
     def initialize_labeling_ui():
         image, metadata, count = get_next_sample()
-        if image is not None:
-            return image, str(metadata["bounds"]), f"Labeled {count} samples."
-        else:
-            return None, "", "No samples loaded."
+        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
@@ -106,32 +91,25 @@ def labeling_ui():
 
 # --- Display UI ---
 def display_ui():
-
     def get_random_cool_images(n):
         cool_samples = [d for d in labeled_data if d["label"] == "cool"]
-        if len(cool_samples) < n:
-          return [Image.frombytes("RGB", (384,384), s["image"]) for s in cool_samples]
-
-        selected_cool = random.sample(cool_samples, n)
-        return [Image.frombytes("RGB", (384,384), s["image"]) for s in selected_cool]
+        return [Image.frombytes("RGB", (384, 384), s["image"]) for s in cool_samples] if len(cool_samples) >= n else []
 
     def get_new_unlabeled_image():
         global data_iter
         try:
             sample = next(data_iter)
             sample = ev.item_to_images(DATASET_SUBSET, sample)
-            image = sample["rgb"][0]
-            metadata = sample["metadata"]
-            return image, str(metadata["bounds"])
+            return sample["rgb"][0], json.dumps(sample["metadata"]["bounds"])
         except StopIteration:
+            print("No more samples in the dataset.")
             return None, None
 
     def refresh_display():
         new_image, new_metadata = get_new_unlabeled_image()
-        if new_image is None:
-            return "No more samples", gr.Image.update(value=None), gr.Gallery.update(value=[])
-
         cool_images = get_random_cool_images(DISPLAY_N_COOL)
+        if new_image is None:
+            return "No more samples", None, []
         return "", new_image, cool_images
 
     with gr.Row():
@@ -139,21 +117,18 @@ def display_ui():
         metadata_display = gr.Textbox(label="Metadata (Bounds)")
 
     with gr.Row():
-        cool_images_gallery = gr.Gallery(
-            label="Cool Examples",
-            value=[],
-            columns=DISPLAY_N_COOL  # Set grid layout here
-        )
-
-    with gr.Row():
-        refresh_button = gr.Button("Refresh")
+        cool_images_gallery = gr.Gallery(label="Cool Examples", value=[], columns=DISPLAY_N_COOL)
 
-    refresh_button.click(fn=refresh_display, inputs=[], outputs=[gr.Textbox(label="Debug"), new_image_component, cool_images_gallery])
+    refresh_button = gr.Button("Refresh")
+    refresh_button.click(
+        fn=refresh_display,
+        inputs=[],
+        outputs=[gr.Textbox(label="Debug"), new_image_component, cool_images_gallery]
+    )
 
-    # Initialize
     def initialize_display_ui():
-      debug, image, gallery = refresh_display()
-      return debug, image, gallery
+        debug, image, gallery = refresh_display()
+        return debug, image, gallery
 
     debug, initial_image, initial_gallery = initialize_display_ui()
     new_image_component.value = initial_image
@@ -162,11 +137,9 @@ def display_ui():
 # --- Main Interface ---
 with gr.Blocks() as demo:
     gr.Markdown("# TerraNomaly")
-
     with gr.Tabs():
         with gr.TabItem("Labeling"):
             labeling_ui()
-
         with gr.TabItem("Display"):
             display_ui()