Spaces:
Running
Running
| from datasets import load_dataset | |
| from functools import partial | |
| from pandas import DataFrame | |
| import earthview as ev | |
| import gradio as gr | |
| import tqdm | |
| import os | |
| DEBUG = "samples" # False, "random", "samples" | |
| if DEBUG == "random": | |
| import numpy as np | |
| def open_dataset(dataset, subset, split, batch_size, shard, only_rgb, state): | |
| nshards = ev.get_nshards(subset) | |
| if shard == -1: | |
| shards = None | |
| else: | |
| shards = [shard] | |
| if DEBUG == "random": | |
| ds = range(batch_size) | |
| elif DEBUG == "samples": | |
| ds = ev.load_parquet(subset, batch_size=batch_size) | |
| elif not DEBUG: | |
| ds = ev.load_dataset(subset, dataset=dataset, split=split, shards=shards, cache_dir="dataset") | |
| dsi = iter(ds) | |
| state["subset"] = subset | |
| state["dsi"] = dsi | |
| return ( | |
| gr.update(label=f"Shard (max {nshards})", value=shard, maximum=nshards), | |
| *get_images(batch_size, only_rgb, state), | |
| state | |
| ) | |
| def get_images(batch_size, only_rgb, state): | |
| subset = state["subset"] | |
| images = [] | |
| metadatas = [] | |
| for i in tqdm.trange(batch_size, desc=f"Getting images"): | |
| if DEBUG == "random": | |
| images.append(np.random.randint(0,255,(384,384,3))) | |
| if not only_rgb: | |
| images.append(np.random.randint(0,255,(100,100,3))) | |
| metadatas.append({"bounds":[[1,1,4,4]], }) | |
| else: | |
| try: | |
| item = next(state["dsi"]) | |
| except StopIteration: | |
| break | |
| metadata = item["metadata"] | |
| item = ev.item_to_images(subset, item) | |
| if subset == "satellogic": | |
| images.extend(item["rgb"]) | |
| if not only_rgb: | |
| images.extend(item["1m"]) | |
| if subset == "sentinel_1": | |
| images.extend(item["10m"]) | |
| if subset == "neon": | |
| images.extend(item["rgb"]) | |
| if not only_rgb: | |
| images.extend(item["chm"]) | |
| images.extend(item["1m"]) | |
| metadatas.append(item["metadata"]) | |
| return images, DataFrame(metadatas) | |
| def update_shape(rows, columns): | |
| return gr.update(rows=rows, columns=columns) | |
| def new_state(): | |
| return gr.State({}) | |
| if __name__ == "__main__": | |
| with gr.Blocks(title="EarthView Viewer", fill_height = True) as demo: | |
| state = new_state() | |
| gr.Markdown(f"# Viewer for [{ev.DATASET}](https://huggingface.co/datasets/satellogic/EarthView) Dataset") | |
| batch_size = gr.Number(10, label = "Batch Size", render=False) | |
| shard = gr.Slider(label="Shard", minimum=0, maximum=10000, step=1, render=False) | |
| table = gr.DataFrame(render = False) | |
| # headers=["Index","TimeStamp","Bounds","CRS"], | |
| gallery = gr.Gallery( | |
| label=ev.DATASET, | |
| interactive=False, | |
| object_fit="scale-down", | |
| columns=5, rows=2, render=False) | |
| with gr.Row(): | |
| dataset = gr.Textbox(label="Dataset", value=ev.DATASET, interactive=False) | |
| subset = gr.Dropdown(choices=ev.get_subsets(), label="Subset", value="satellogic", ) | |
| split = gr.Textbox(label="Split", value="train") | |
| initial_shard = gr.Number(label = "Initial shard", value=10, info="-1 for whole dataset") | |
| only_rgb = gr.Checkbox(label="Only RGB", value=True) | |
| gr.Button("Load (minutes)").click( | |
| open_dataset, | |
| inputs=[dataset, subset, split, batch_size, initial_shard, only_rgb, state], | |
| outputs=[shard, gallery, table, state]) | |
| gallery.render() | |
| with gr.Row(): | |
| batch_size.render() | |
| rows = gr.Number(2, label="Rows") | |
| columns = gr.Number(5, label="Coluns") | |
| rows.change(update_shape, [rows, columns], [gallery]) | |
| columns.change(update_shape, [rows, columns], [gallery]) | |
| with gr.Row(): | |
| shard.render() | |
| shard.release( | |
| open_dataset, | |
| inputs=[dataset, subset, split, batch_size, shard, only_rgb, state], | |
| outputs=[shard, gallery, table, state]) | |
| btn = gr.Button("Next Batch (same shard)", scale=0) | |
| btn.click(get_images, [batch_size, only_rgb, state], [gallery, table]) | |
| btn.click() | |
| table.render() | |
| demo.launch(show_api=False) | |