Update app.py

app.py

@@ -1,5 +1,8 @@
 import gradio as gr
 import earthview as ev
+from PIL import Image
+import numpy as np
+import random
 import os
 import json
 import utils
@@ -7,8 +10,11 @@ from pandas import DataFrame
 
 # --- Configuration ---
 DATASET_SUBSET = "satellogic"
+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])
@@ -35,14 +41,30 @@ def get_next_sample():
         print("No more samples in the dataset.")
         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(label, state):
     global labeled_data
     sample = state["sample"]
     if sample is None:
-        return "No image to label", DataFrame()
+        return "No image to label", None, DataFrame()
 
-    image_bytes = sample["rgb"][0].convert("RGB").tobytes()  # Convert PIL Image to bytes
+    image = sample["rgb"][0]
+    image_bytes = image.convert("RGB").tobytes()
     labeled_data.append({
         "image": image_bytes,
         "metadata": sample["metadata"],
@@ -55,45 +77,91 @@ def save_labeled_data(label, state):
     new_sample = get_next_sample()
     if new_sample is None:
         state["sample"] = None
-        return "Dataset exhausted.", DataFrame()
+        return "Dataset exhausted.", [], 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 "", DataFrame([new_metadata])
+    return "", [new_image], DataFrame([new_metadata])  # Return a list of image data
 
 # --- Gradio Interface ---
+# --- Labeling UI ---
 def labeling_ui():
     state = gr.State({"sample": None, "subset": DATASET_SUBSET})
 
     with gr.Row():
-        cool_button = gr.Button("Cool")
-        not_cool_button = gr.Button("Not Cool")
+        with gr.Column():
+            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()
+                images, metadata = get_images(1, {"sample": None, "subset": DATASET_SUBSET})
+                return sample, images, metadata
+
+            initial_sample, initial_images, initial_metadata = initialize_labeling_ui()
+            gallery.value = initial_images  # Initialize with a list of images
+            table.value = initial_metadata
+            state.value = {"sample": initial_sample, "subset": DATASET_SUBSET}
+
+            cool_button.click(
+                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=lambda label, state: save_labeled_data(label, state),
+                inputs=[gr.Textbox(visible=False, value="not cool"), state],
+                outputs=[gr.Textbox(label="Debug"), gallery, table]
+            )
+
+# --- 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 []
+
+    def get_new_unlabeled_image():
+        global data_iter
+        try:
+            sample = next(data_iter)
+            sample = ev.item_to_images(DATASET_SUBSET, sample)
+            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()
+        cool_images = get_random_cool_images(DISPLAY_N_COOL)
+        if new_image is None:
+            return "No more samples", None, []
+        return "", new_image, cool_images
 
-    table = gr.DataFrame(datatype="html")
+    with gr.Row():
+        new_image_component = gr.Image(label="New Image", type="pil")
+        metadata_display = gr.Textbox(label="Metadata (Bounds)")
+    with gr.Row():
+        cool_images_gallery = gr.Gallery(label="Cool Examples", value=[], columns=DISPLAY_N_COOL)
 
-    def initialize_labeling_ui():
-        sample = get_next_sample()
-        if sample is None:
-            return {"sample": None, "subset": DATASET_SUBSET}, DataFrame()
-        metadata = sample["metadata"]
-        metadata["map"] = f'<a href="{utils.get_google_map_link(sample, DATASET_SUBSET)}" target="_blank">🧭</a>'
-        return {"sample": sample, "subset": DATASET_SUBSET}, DataFrame([metadata])
+    refresh_button = gr.Button("Refresh")
+    refresh_button.click(
+        fn=refresh_display,
+        inputs=[],
+        outputs=[gr.Textbox(label="Debug"), new_image_component, cool_images_gallery]
+    )
 
-    initial_state, initial_metadata = initialize_labeling_ui()
-    table.value = initial_metadata
-    state.value = initial_state
+    def initialize_display_ui():
+        debug, image, gallery = refresh_display()
+        return debug, image, gallery
 
-    cool_button.click(
-        fn=lambda label, state: save_labeled_data(label, state),
-        inputs=[gr.Textbox(visible=False, value="cool"), state],
-        outputs=[gr.Textbox(label="Debug"), table]
-    )
-    not_cool_button.click(
-        fn=lambda label, state: save_labeled_data(label, state),
-        inputs=[gr.Textbox(visible=False, value="not cool"), state],
-        outputs=[gr.Textbox(label="Debug"), table]
-    )
+    debug, initial_image, initial_gallery = initialize_display_ui()
+    new_image_component.value = initial_image
+    cool_images_gallery.value = initial_gallery
 
 # --- Main Interface ---
 with gr.Blocks() as demo:
@@ -101,5 +169,7 @@ with gr.Blocks() as demo:
     with gr.Tabs():
         with gr.TabItem("Labeling"):
             labeling_ui()
+        with gr.TabItem("Display"):
+            display_ui()
 
 demo.launch(debug=True)