Update app.py

app.py

@@ -4,61 +4,77 @@ import utils
 import random
 import pandas as pd
 import os
+from itertools import islice
 
-# Load the Satellogic dataset
-dataset = ev.load_dataset("satellogic", streaming=True).shuffle(seed=42)
-data_iter = iter(dataset)
-
-# File to store labels (will create if it doesn't exist)
+# Configuration
+chunk_size = 100  # Size of the chunks to shuffle
 label_file = "labels.csv"
 
-# Initialize a DataFrame to hold labels (or load existing)
+# Load the Satellogic dataset (streaming)
+dataset = ev.load_dataset("satellogic", streaming=True)
+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"])
 
 def get_next_image():
-    global data_iter, labels_df
+    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
+        if not shuffled_chunk or chunk_iter is None:
+            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))
+                chunk = list(islice(data_iter, chunk_size))
+                if not chunk:
+                    print("Still no data after reset.")
+                    return None, "Dataset exhausted", None, None
+
+            random.shuffle(chunk)
+            shuffled_chunk = chunk
+            chunk_iter = iter(shuffled_chunk)
 
-    while True:  # Keep iterating until we find an unlabeled image
         try:
-            sample = next(data_iter)
+            sample = next(chunk_iter)
+            sample = ev.item_to_images("satellogic", sample)
+            image = sample["rgb"][0]
+            metadata = sample["metadata"]
+            bounds = metadata["bounds"]
+            google_maps_link = utils.get_google_map_link(sample, "satellogic")
+            image_id = str(bounds)
+
+            if image_id not in labels_df["image_id"].values:
+                return image, image_id, bounds, google_maps_link
         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
+            # 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_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]})
+    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
+# Gradio interface (no changes needed here)
 iface = gr.Interface(
     fn=rate_image,
     inputs=[
@@ -78,10 +94,8 @@ iface = gr.Interface(
     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=[
@@ -90,5 +104,4 @@ iface.launch(
         initial_bounds,
         initial_google_maps_link,
     ],
-
 )