Update app.py

app.py

@@ -1,19 +1,8 @@
-import os
 import pandas as pd
 from PIL import Image
-from itertools import islice
-import earthview as ev  # Your custom library
+import earthview as ev
 import gradio as gr
-
-# Path to the first shard of the Satellogic dataset
-filename = os.path.join("dataset", "satellogic", "train-00000-of-00065.parquet")
-
-# Check if the file exists
-if not os.path.isfile(filename):
-    raise FileNotFoundError(f"{filename} not found")
-
-# Loads a dataset with pandas, this loads a single file
-data = pd.read_parquet(filename)
+from datasets import load_dataset
 
 # Global variables to manage state
 current_image = None
@@ -34,22 +23,26 @@ def item_to_bounds_timestamps(sample):
         timestamps.append(t)
     return bounds, timestamps
 
+# Load the dataset directly from Hugging Face
+data = load_dataset("satellogic/EarthView", "satellogic", split="train", streaming=True)
+
 # Function to load and display the next image
 def load_next_image():
     global current_image, image_id, current_metadata_index, bounds, timestamps
 
-    if current_metadata_index >= len(data):
-        return None, "No more images", None
+    try:
+        sample = next(iter(data))
+        bounds_sample, timestamps_sample = item_to_bounds_timestamps(sample)
 
-    sample = data.iloc[current_metadata_index]
-    bounds_sample, timestamps_sample = item_to_bounds_timestamps(sample)
+        # Use earthview library to convert arrays to PIL images
+        sample = ev.item_to_images("satellogic", sample)
+        current_image = sample["rgb"][0]  # Get the first image
+        image_id += 1
 
-    # Use earthview library to convert arrays to PIL images
-    sample = ev.item_to_images("satellogic", sample.to_dict())
-    current_image = sample["rgb"][0]  # Get the first image
-    image_id += 1
+        return current_image, f"Image ID: {image_id}", bounds_sample[0]
 
-    return current_image, f"Image ID: {image_id}", bounds_sample[0]
+    except StopIteration:
+        return None, "No more images", None
 
 # Function to handle rating submission
 def submit_rating(rating, bounds_str):
@@ -58,7 +51,7 @@ def submit_rating(rating, bounds_str):
     ratings.append(rating)
     bounds.append(bounds_str)
     timestamps.append("timestamp")  # Use a valid timestamp if available
-    
+
     current_metadata_index += 1
 
     return load_next_image()