Update app.py

app.py

@@ -1,74 +1,145 @@
-import pandas as pd
-from PIL import Image
-import earthview as ev
 import gradio as gr
-from datasets import load_dataset
-
-# Global variables to manage state
-current_image = None
-current_metadata_index = 0
-ratings = []
-image_id = 0
-bounds = []
-timestamps = []
-
-# Function to transform a metadata sample to bounds and timestamps
-def item_to_bounds_timestamps(sample):
-    bounds_list = sample["metadata"]["bounds"]
-    timestamp_list = sample["metadata"]["timestamp"]
-    bounds = []
-    timestamps = []
-    for b, t in zip(bounds_list, timestamp_list):
-        bounds.append(b)
-        timestamps.append(t)
-    return bounds, timestamps
-
-# Load the dataset directly from Hugging Face
-data = load_dataset("satellogic/EarthView", "satellogic", split="train", streaming=True)
-
-# Function to load and display the next image
-def load_next_image():
-    global current_image, image_id, current_metadata_index, bounds, timestamps
+import earthview as ev
+from PIL import Image
+import numpy as np
+import random
+import os
+import json
 
-    try:
-        sample = next(iter(data))
-        bounds_sample, timestamps_sample = item_to_bounds_timestamps(sample)
+# --- 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.
+# --- Load Dataset ---
+dataset = ev.load_dataset(DATASET_SUBSET, shards=[SAMPLE_SEED])
+data_iter = iter(dataset)
+
+# --- Load Labeled Data (if it exists) ---
+def load_labeled_data():
+    if os.path.exists(LABELED_DATA_FILE):
+        with open(LABELED_DATA_FILE, "r") as f:
+            return json.load(f)
+    else:
+        return []
 
-        # Use earthview library to convert arrays to PIL images
-        sample = ev.item_to_images("satellogic", sample)
-        current_image = sample["rgb"][0]  # Get the first image
-        image_id += 1
+labeled_data = load_labeled_data()
 
-        return current_image, f"Image ID: {image_id}", bounds_sample[0]
+# --- Get Next Sample for Labeling ---
+def get_next_sample():
+    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, metadata, len(labeled_data)
     except StopIteration:
-        return None, "No more images", None
+        return None, None, None
 
-# Function to handle rating submission
-def submit_rating(rating, bounds_str):
-    global image_id, current_metadata_index, ratings, bounds, timestamps
+# --- Save Labeled Data ---
+def save_labeled_data(image, metadata, label):
+    global labeled_data
+    labeled_data.append({
+        "image": image.convert("RGB").tobytes(),  # Convert to RGB and then serialize
+        "metadata": metadata,
+        "label": label
+    })
 
-    ratings.append(rating)
-    bounds.append(bounds_str)
-    timestamps.append("timestamp")  # Use a valid timestamp if available
+    with open(LABELED_DATA_FILE, "w") as f:
+        json.dump(labeled_data, f)
 
-    current_metadata_index += 1
+    image, metadata, count = get_next_sample()
 
-    return load_next_image()
+    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."
+
+# --- Gradio Interface ---
+
+# --- Labeling UI ---
+def labeling_ui():
+  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")
+
+          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"), 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"), inputs=[image_component, metadata_text], outputs=[gr.Textbox(label="Debug"), image_component, metadata_text, label_count_text])
+
+  # Initialize with the first sample
+  image, metadata, count = get_next_sample()
+  if image is not None:
+      image_component.value = image
+      metadata_text.value = str(metadata["bounds"])
+      label_count_text.value = f"Labeled {count} samples."
+
+# --- 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]
+
+    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"])
+        except StopIteration:
+            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)
+        return "", new_image, cool_images
+
+    with gr.Row():
+        new_image_component = gr.Image(label="New Image", type="pil")
+        metadata_display = gr.Textbox(label="Metadata (Bounds)")
 
-# Gradio Interface Layout
-with gr.Blocks() as demo:
     with gr.Row():
-        with gr.Column():
-            image_display = gr.Image(label="Satellite Image", type="pil")
-            image_id_display = gr.Textbox(label="Image ID")
-            bounds_display = gr.Textbox(label="Bounds")
-            load_button = gr.Button("Load Next Image")
-            load_button.click(fn=load_next_image, outputs=[image_display, image_id_display, bounds_display])
-        with gr.Column():
-            rating_radio = gr.Radio(["0", "1"], label="Rating (0 = No, 1 = Yes)")
-            submit_button = gr.Button("Submit Rating")
-            submit_button.click(fn=lambda rating, bounds_str: submit_rating(rating, bounds_str), inputs=[rating_radio, bounds_display], outputs=[image_display, image_id_display, bounds_display])
-
-# Launch the Gradio interface
-demo.launch()
+        cool_images_gallery = gr.Gallery(label="Cool Examples", value=[])
+        cool_images_gallery.style(grid=DISPLAY_N_COOL)
+
+    with gr.Row():
+        refresh_button = gr.Button("Refresh")
+
+    refresh_button.click(fn=refresh_display, inputs=[], outputs=[gr.Textbox(label="Debug"), new_image_component, cool_images_gallery])
+
+    # Initialize
+    debug, image, gallery = refresh_display()
+    new_image_component.value = image
+    cool_images_gallery.value = gallery
+
+# --- 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()
+
+demo.launch(debug=True)