Update app.py

app.py

@@ -5,8 +5,6 @@ import random
 import pandas as pd
 import os
 from itertools import islice
-import pyarrow as pa
-import pyarrow.parquet as pq
 
 # Configuration
 chunk_size = 100  # Size of the chunks to shuffle
@@ -18,7 +16,7 @@ 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 (make it global for the download function)
+# Initialize or load labels DataFrame
 labels_df = None
 if os.path.exists(label_file):
     labels_df = pd.read_csv(label_file)
@@ -34,7 +32,7 @@ def get_next_image():
             chunk = list(islice(data_iter, chunk_size))
             if not chunk:  # If the dataset is exhausted, reset the iterator
                 print("Dataset exhausted, resetting iterator.")
-                data_iter = iter(ev.load_dataset("satellogic", streaming=True))
+                reset_dataset_iterator() # Use the reset function
                 chunk = list(islice(data_iter, chunk_size))
                 if not chunk:
                     print("Still no data after reset.")
@@ -53,7 +51,9 @@ 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:
+            if labels_df is not None and image_id not in labels_df["image_id"].values:
+                return image, image_id, bounds, google_maps_link
+            elif labels_df is None: # Handle case where labels_df is not initialized yet
                 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
@@ -86,6 +86,17 @@ def save_labels_parquet():
     else:
         return None
 
+def reset_dataset_iterator():
+    global data_iter, shuffled_chunk, chunk_iter
+    data_iter = iter(ev.load_dataset("satellogic", streaming=True))
+    shuffled_chunk = []
+    chunk_iter = None
+
+def load_different_batch():
+    print("Loading a different batch of images...")
+    reset_dataset_iterator()
+    return get_next_image() # Return the first image from the new batch
+
 # Gradio interface
 with gr.Blocks() as iface:
     image_out = gr.Image(label="Satellite Image")
@@ -95,6 +106,7 @@ with gr.Blocks() as iface:
     rating_radio = gr.Radio(["Cool", "Not Cool"], label="Rating")
     with gr.Row():
         submit_button = gr.Button("Submit Rating")
+        different_batch_button = gr.Button("Load Different Batch") # New button
         download_button = gr.Button("Download Labels (Parquet)")
     download_output = gr.File(label="Download Labeled Data")
 
@@ -110,6 +122,12 @@ with gr.Blocks() as iface:
         outputs=[download_output],
     )
 
+    different_batch_button.click(
+        fn=load_different_batch,
+        inputs=[],
+        outputs=[image_out, image_id_out, bounds_out, google_maps_link_out]
+    )
+
      # Load the first image
     initial_image, initial_image_id, initial_bounds, initial_google_maps_link = get_next_image()