Update app.py

app.py

@@ -10,6 +10,7 @@ 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)
@@ -17,11 +18,15 @@ data_iter = iter(dataset)
 shuffled_chunk = []  # Initialize an empty list to hold the current chunk
 chunk_iter = None  # Initialize the chunk iterator
 
-# Initialize an empty list to hold labels temporarily
-labels_list = []
+# Initialize or load labels DataFrame (make it global for the download function)
+labels_df = None
+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_list, shuffled_chunk, chunk_iter
+    global data_iter, labels_df, shuffled_chunk, chunk_iter
 
     while True:
         # If we don't have a current chunk or it's exhausted, get a new one
@@ -48,71 +53,70 @@ def get_next_image():
             google_maps_link = utils.get_google_map_link(sample, "satellogic")
             image_id = str(bounds)
 
-            # 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):
+            if image_id not in labels_df["image_id"].values:
                 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
             shuffled_chunk = []
             chunk_iter = None
 
-def rate_image(image_id, bounds, rating, google_maps_link):
-    global labels_list
+def rate_image(image_id, bounds, rating):
+    global labels_df
 
-    labels_list.append(
+    new_row = pd.DataFrame(
         {
-            "image_id": image_id,
-            "bounds": bounds,
-            "rating": rating,
-            "google_maps_link": google_maps_link,
+            "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.
         }
     )
+    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
-    if labels_list:
-       table = pa.Table.from_pylist(labels_list)
-       pq.write_table(table, "labeled_data.parquet")
-       return "labeled_data.parquet"
+def save_labels_parquet():
+    global labels_df
+    if labels_df is not None and not labels_df.empty:
+        table = pa.Table.from_pandas(labels_df)
+        pq.write_table(table, 'labeled_data.parquet')
+        return 'labeled_data.parquet'
     else:
         return None
 
 # Gradio interface
 with gr.Blocks() as iface:
+    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)
+    rating_radio = gr.Radio(["Cool", "Not Cool"], label="Rating")
     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")
-            with gr.Row():
-                submit_button = gr.Button("Submit Rating")
-                download_button = gr.Button("Download Labels")
-            download_file = gr.File(label="Download")
+        submit_button = gr.Button("Submit Rating")
+        download_button = gr.Button("Download Labels (Parquet)")
+    download_output = gr.File(label="Download Labeled Data")
 
     submit_button.click(
-        rate_image,
-        inputs=[image_id_out, bounds_out, rating_radio, google_maps_link_out],
-        outputs=[image_out, image_id_out, bounds_out, google_maps_link_out]
+        fn=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]
+        fn=save_labels_parquet,
+        inputs=[],
+        outputs=[download_output],
     )
 
-    # Get the first image and its details
+     # Load the first image
     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
+    # Set initial values
+    if initial_image:
+        iface.load(lambda: (initial_image, initial_image_id, initial_bounds, initial_google_maps_link),
+                   inputs=None,
+                   outputs=[image_out, image_id_out, bounds_out, google_maps_link_out])
 
 iface.launch(share=True)