app.py

import pandas as pd
from PIL import Image
import earthview as ev
import gradio as gr
from datasets import load_dataset

# Global variables to manage state
current_image = None
current_metadata_index = 0
ratings = []
image_id = 0
bounds = []
timestamps = []

# Function to transform a metadata sample to bounds and timestamps
def item_to_bounds_timestamps(sample):
    bounds_list = sample["metadata"]["bounds"]
    timestamp_list = sample["metadata"]["timestamp"]
    bounds = []
    timestamps = []
    for b, t in zip(bounds_list, timestamp_list):
        bounds.append(b)
        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

    try:
        sample = next(iter(data))
        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

        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):
    global image_id, current_metadata_index, ratings, bounds, timestamps

    ratings.append(rating)
    bounds.append(bounds_str)
    timestamps.append("timestamp")  # Use a valid timestamp if available

    current_metadata_index += 1

    return load_next_image()

# Gradio Interface Layout
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            image_display = gr.Image(label="Satellite Image", type="pil")
            image_id_display = gr.Textbox(label="Image ID")
            bounds_display = gr.Textbox(label="Bounds")
            load_button = gr.Button("Load Next Image")
            load_button.click(fn=load_next_image, outputs=[image_display, image_id_display, bounds_display])
        with gr.Column():
            rating_radio = gr.Radio(["0", "1"], label="Rating (0 = No, 1 = Yes)")
            submit_button = gr.Button("Submit Rating")
            submit_button.click(fn=lambda rating, bounds_str: submit_rating(rating, bounds_str), inputs=[rating_radio, bounds_display], outputs=[image_display, image_id_display, bounds_display])

# Launch the Gradio interface
demo.launch()