Update app.py

app.py

@@ -1,123 +1,94 @@
 import gradio as gr
 import earthview as ev
-from PIL import Image
-import numpy as np
+import utils
 import random
+import pandas as pd
 import os
-import json
-import utils
-from pandas import DataFrame
-
-# --- Configuration ---
-DATASET_SUBSET = "satellogic"
-LABELED_DATA_FILE = "labeled_data.json"
-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])
+# Load the Satellogic dataset
+dataset = ev.load_dataset("satellogic", streaming=True).shuffle(seed=42)
 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 []
-
-labeled_data = load_labeled_data()
-
-# --- 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)
-        return sample
-    except StopIteration:
-        print("No more samples in the dataset.")
-        return None
-
-# --- Save Labeled Data ---
-def save_labeled_data(image, label, state):
-    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
-
-    labeled_data.append({
-        "image": image_bytes,
-        "metadata": state["metadata"],
-        "label": label
-    })
-
-    with open(LABELED_DATA_FILE, "w") as f:
-        json.dump(labeled_data, f)
-
-    new_sample = get_next_sample()
-    if new_sample is None:
-        print("Dataset exhausted.")
-        return "Dataset exhausted.", None, DataFrame()
-
-    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>'
-    state["metadata"] = new_metadata  # Update metadata in state
-
-    # Convert new PIL Image to bytes for sending to client-side
-    if new_image is not None:
-        new_image_bytes = new_image.convert("RGB").tobytes()
-    else:
-        new_image_bytes = None
-
-    return "", new_image_bytes, DataFrame([new_metadata])
-
-# --- Gradio Interface ---
-# --- Labeling UI ---
-def labeling_ui():
-    state = gr.State({"sample": None, "metadata": None, "subset": DATASET_SUBSET})
-
-    with gr.Row():
-        with gr.Column():
-            image_component = gr.Image(label="Satellite Image", type="pil", interactive=False)
-            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()
-                if sample is None:
-                    return {"sample": None, "metadata": None, "subset": DATASET_SUBSET}, None, DataFrame()
-                image = sample["rgb"][0]
-                metadata = sample["metadata"]
-                metadata["map"] = f'<a href="{utils.get_google_map_link(sample, DATASET_SUBSET)}" target="_blank">🧭</a>'
-                return {"sample": sample, "metadata": metadata, "subset": DATASET_SUBSET}, image, DataFrame([metadata])
-
-            initial_state, initial_image, initial_metadata = initialize_labeling_ui()
-            image_component.value = initial_image
-            table.value = initial_metadata
-            state.value = initial_state
-
-            cool_button.click(
-                fn=lambda image, label, state: save_labeled_data(image, label, state),
-                inputs=[image_component, gr.Textbox(visible=False, value="cool"), state],
-                outputs=[gr.Textbox(label="Debug"), image_component, table]
-            )
-            not_cool_button.click(
-                fn=lambda image, label, state: save_labeled_data(image, label, state),
-                inputs=[image_component, gr.Textbox(visible=False, value="not cool"), state],
-                outputs=[gr.Textbox(label="Debug"), image_component, table]
-            )
-
-# --- Main Interface ---
-with gr.Blocks() as demo:
-    gr.Markdown("# TerraNomaly")
-    with gr.Tabs():
-        with gr.TabItem("Labeling"):
-            labeling_ui()
-
-demo.launch(debug=True)
+# File to store labels (will create if it doesn't exist)
+label_file = "labels.csv"
+
+# Initialize a DataFrame to hold labels (or load existing)
+if os.path.exists(label_file):
+    labels_df = pd.read_csv(label_file)
+else:
+    labels_df = pd.DataFrame(columns=["image_id", "bounds", "rating", "google_maps_link"])
+
+def get_next_image():
+    global data_iter, labels_df
+
+    while True:  # Keep iterating until we find an unlabeled image
+        try:
+            sample = next(data_iter)
+        except StopIteration:
+            #refresh the dataset if we reach the end
+            dataset = ev.load_dataset("satellogic", streaming=True).shuffle(seed=random.randint(0, 1000000))
+            data_iter = iter(dataset)
+            continue
+
+        sample = ev.item_to_images("satellogic", sample)
+        image = sample["rgb"][0]  # Get the first RGB image
+        metadata = sample["metadata"]
+
+        bounds = metadata["bounds"]
+        google_maps_link = utils.get_google_map_link(sample, "satellogic")
+        #generate a unique image ID:
+        image_id = (str(bounds))
+
+        # Check if image is already labeled
+        if image_id not in labels_df["image_id"].values:
+            return image, image_id, bounds, google_maps_link
+
+def rate_image(image_id, bounds, rating, google_maps_link):
+    global labels_df
+
+    # Add the rating to the DataFrame
+    new_row = pd.DataFrame({"image_id": [image_id], "bounds": [bounds], "rating": [rating], "google_maps_link": [google_maps_link]})
+    labels_df = pd.concat([labels_df, new_row], ignore_index=True)
+
+    # Save the DataFrame to CSV
+    labels_df.to_csv(label_file, index=False)
+
+    # Get the next image and its details
+    next_image, next_image_id, next_bounds, next_google_maps_link = get_next_image()
+
+    return next_image, next_image_id, next_bounds, next_google_maps_link
+
+# Define the Gradio interface
+iface = gr.Interface(
+    fn=rate_image,
+    inputs=[
+        gr.Textbox(label="Image ID", visible=False),
+        gr.Textbox(label="Bounds", visible=False),
+        gr.Radio(["Cool", "Not Cool"], label="Rating"),
+        gr.Textbox(label="Google Maps Link"),
+    ],
+    outputs=[
+        gr.Image(label="Satellite Image"),
+        gr.Textbox(label="Image ID", visible=False),
+        gr.Textbox(label="Bounds", visible=False),
+        gr.Textbox(label="Google Maps Link"),
+    ],
+    title="TerraNomaly - Satellite Image Labeling",
+    description="Rate satellite images as 'Cool' or 'Not Cool'.",
+    live=False,
+)
+
+# Get the first image and its details
+initial_image, initial_image_id, initial_bounds, initial_google_maps_link = get_next_image()
+
+# Set the initial values for the output components
+iface.launch(
+    share=True,
+    initial_outputs=[
+        initial_image,
+        initial_image_id,
+        initial_bounds,
+        initial_google_maps_link,
+    ],
+
+)