Update app.py

app.py

@@ -5,10 +5,12 @@ import random
 import pandas as pd
 import os
 from itertools import islice
+import json
+import pyarrow as pa
+import pyarrow.parquet as pq
 
 # Configuration
 chunk_size = 100  # Size of the chunks to shuffle
-label_file = os.path.join(os.path.dirname(__file__), "labels.csv")  # Save CSV in the same directory as the script
 
 # Load the Satellogic dataset (streaming)
 dataset = ev.load_dataset("satellogic", streaming=True)
@@ -16,14 +18,11 @@ data_iter = iter(dataset)
 shuffled_chunk = []  # Initialize an empty list to hold the current chunk
 chunk_iter = None  # Initialize the chunk iterator
 
-# Initialize or load labels DataFrame
-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"])
+# Initialize an empty list to hold labels temporarily
+labels_list = []
 
 def get_next_image():
-    global data_iter, labels_df, shuffled_chunk, chunk_iter
+    global data_iter, labels_list, shuffled_chunk, chunk_iter
 
     while True:
         # If we don't have a current chunk or it's exhausted, get a new one
@@ -50,7 +49,8 @@ def get_next_image():
             google_maps_link = utils.get_google_map_link(sample, "satellogic")
             image_id = str(bounds)
 
-            if image_id not in labels_df["image_id"].values:
+            # Check if this image has already been labeled (based on image_id)
+            if not any(label["image_id"] == image_id for label in labels_list):
                 return image, image_id, bounds, google_maps_link
         except StopIteration:
             # Current chunk is exhausted, reset chunk variables to get a new one in the next iteration
@@ -58,42 +58,58 @@ def get_next_image():
             chunk_iter = None
 
 def rate_image(image_id, bounds, rating):
-    global labels_df
+    global labels_list
 
-    new_row = pd.DataFrame(
+    labels_list.append(
         {
-            "image_id": [image_id],
-            "bounds": [bounds],
-            "rating": [rating],
-            "google_maps_link": [""], # this isn't necessary to pass to the function since we aren't updating it here.
+            "image_id": image_id,
+            "bounds": bounds,
+            "rating": rating,
+            "google_maps_link": "", # Adding google maps link to the data to be downloaded
         }
     )
-    labels_df = pd.concat([labels_df, new_row], ignore_index=True)
-    labels_df.to_csv(label_file, index=False)
 
     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
 
+def save_labels():
+    global labels_list
+    table = pa.Table.from_pylist(labels_list)
+    pq.write_table(table, "labeled_data.parquet")
+    return "labeled_data.parquet"
+
 # 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"), # Remove google maps link as an input
-    ],
-    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", visible=True), # Add google maps link as an output
-    ],
-    title="TerraNomaly - Satellite Image Labeling",
-    description="Rate satellite images as 'Cool' or 'Not Cool'.",
-    live=False,
-)
-
-iface.launch(
-    share=True
-)
+with gr.Blocks() as iface:
+    with gr.Row():
+        with gr.Column():
+            image_out = gr.Image(label="Satellite Image")
+            image_id_out = gr.Textbox(label="Image ID", visible=False)
+            bounds_out = gr.Textbox(label="Bounds", visible=False)
+            google_maps_link_out = gr.Textbox(label="Google Maps Link", visible=True)
+        with gr.Column():
+            rating_radio = gr.Radio(["Cool", "Not Cool"], label="Rating")
+            submit_button = gr.Button("Submit Rating")
+            download_button = gr.Button("Download Labels")
+            download_file = gr.File(label="Download")
+
+    submit_button.click(
+        rate_image,
+        inputs=[image_id_out, bounds_out, rating_radio],
+        outputs=[image_out, image_id_out, bounds_out, google_maps_link_out]
+    )
+
+    download_button.click(
+        save_labels,
+        inputs=None,
+        outputs=[download_file]
+    )
+
+    # Get the first image and its details
+    initial_image, initial_image_id, initial_bounds, initial_google_maps_link = get_next_image()
+
+    image_out.value = initial_image
+    image_id_out.value = initial_image_id
+    bounds_out.value = initial_bounds
+    google_maps_link_out.value = initial_google_maps_link
+
+iface.launch(share=True)