moved code to earthview.py (and other things)

app.py

@@ -1,4 +1,4 @@
-from datasets import load_dataset, get_dataset_config_names
+from datasets import load_dataset
 from functools import partial
 from pandas import DataFrame
 import earthview as ev
@@ -6,73 +6,70 @@ import gradio as gr
 import tqdm
 import os
 
-DEBUG = False
+DEBUG = "samples"       # False, "random", "samples"
 
-if DEBUG:
+if DEBUG == "random":
     import numpy as np
 
-def open_dataset(dataset, set_name, split, batch_size, state, shard = -1):
-    if shard == -1:
-        # Trick to open the whole dataset
-        data_files = None
-        shards = 100
-    else:
-        config = ev.sets[set_name].get("config", set_name)
-        shards = ev.sets[set_name]["shards"]
-        path   = ev.sets[set_name].get("path", set_name)
-        data_files = {"train":[f"{path}/{split}-{shard:05d}-of-{shards:05d}.parquet"]}
+def open_dataset(dataset, subset, split, batch_size, shard, only_rgb, state):
 
-    if DEBUG:
-        dsi = range(100)
+    nshards = ev.get_nshards(subset)
+    
+    if shard == -1:
+        shards = None
     else:
-        ds = load_dataset(
-            dataset,
-            config,
-            split=split,
-            cache_dir="dataset",
-            data_files=data_files,
-            streaming=True,
-            token=os.environ.get("HF_TOKEN", None))
+        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)
+    dsi = iter(ds)
 
-    state["config"]  = config
+    state["subset"]  = subset
     state["dsi"] = dsi
     return (
-        gr.update(label=f"Shards (max {shards})", value=shard, maximum=shards),
-        *get_images(batch_size, state),
+        gr.update(label=f"Shard (max {nshards})", value=shard, maximum=nshards),
+        *get_images(batch_size, only_rgb, state),
         state
     )
     
-def get_images(batch_size, state):
-    config = state["config"]
+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:
+        if DEBUG == "random":
             images.append(np.random.randint(0,255,(384,384,3)))
-            images.append(np.random.randint(0,255,(100,100,3)))
+            if not only_rgb:
+                images.append(np.random.randint(0,255,(100,100,3)))
 
-            metadata = {"bounds":[[1,1,4,4]], }
+            metadatas.append({"bounds":[[1,1,4,4]], })
         else:
             try:
                 item = next(state["dsi"])
             except StopIteration:
                 break
             metadata = item["metadata"]
-            item = ev.item_to_images(config, item)
+            item = ev.item_to_images(subset, item)
 
-            if  config == "satellogic":
+            if  subset == "satellogic":
                 images.extend(item["rgb"])
-                # images.extend(item["1m"])
-            if  config == "sentinel_1":
+                if not only_rgb:
+                    images.extend(item["1m"])
+            if  subset == "sentinel_1":
                 images.extend(item["10m"])
-            if  config == "default":
+            if  subset == "neon":
                 images.extend(item["rgb"])
-                images.extend(item["chm"])
-                images.extend(item["1m"])
+                if not only_rgb:
+                    images.extend(item["chm"])
+                    images.extend(item["1m"])
+            
             metadatas.append(item["metadata"])
 
     return images, DataFrame(metadatas)
@@ -84,7 +81,7 @@ def new_state():
     return gr.State({})
 
 if __name__ == "__main__":
-    with gr.Blocks(title="Dataset Explorer", fill_height = True) as demo:
+    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")
@@ -101,13 +98,14 @@ if __name__ == "__main__":
 
         with gr.Row():
             dataset = gr.Textbox(label="Dataset", value=ev.DATASET, interactive=False)
-            config = gr.Dropdown(choices=ev.get_sets(), label="Config", value="satellogic", )
+            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, config, split, batch_size, state, initial_shard],
+                inputs=[dataset, subset, split, batch_size, initial_shard, only_rgb, state],
                 outputs=[shard, gallery, table, state])
 
         gallery.render()
@@ -125,11 +123,11 @@ if __name__ == "__main__":
             shard.render()
             shard.release(
                 open_dataset,
-                inputs=[dataset, config, split, batch_size, state, shard],
+                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, state], [gallery, table])
+            btn.click(get_images, [batch_size, only_rgb, state], [gallery, table])
             btn.click()
         
         table.render()

earthview.py

@@ -1,7 +1,13 @@
+from datasets import load_dataset as _load_dataset
+from os import environ
 from PIL import Image
 import numpy as np
 import json
 
+from pyarrow.parquet import ParquetFile
+from pyarrow import Table as pa_Table
+from datasets import Dataset
+
 DATASET = "satellogic/EarthView"
 
 sets = {
@@ -18,10 +24,58 @@ sets = {
     }
 }
 
-def get_sets():
+def get_subsets():
     return sets.keys()
 
-def item_to_images(config, item):
+def get_nshards(subset):
+    return sets[subset]["shards"]
+
+def get_path(subset):
+    return sets[subset].get("path", subset)
+
+def get_config(subset):
+    return sets[subset].get("config", subset)
+
+def load_dataset(subset, dataset="satellogic/EarthView", split="train", shards = None, streaming=True, **kwargs):
+    config = get_config(subset)
+    nshards = get_nshards(subset)
+    path   = get_path(subset)
+    if shards is None:
+        data_files = None
+    else:
+        data_files = [f"{path}/{split}-{shard:05d}-of-{nshards:05d}.parquet" for shard in shards]
+        data_files = {split: data_files}
+
+    ds = _load_dataset(
+        path=dataset,
+        name=config,
+        save_infos=True,
+        split=split,
+        data_files=data_files,
+        streaming=streaming,
+        token=environ.get("HF_TOKEN", None),
+        **kwargs)
+
+    return ds    
+
+def load_parquet(subset_or_filename, batch_size=100):
+    if subset_or_filename in get_subsets():
+        pqfile = ParquetFile(f"dataset/{subset_or_filename}/sample.parquet")
+    else:
+        pqfile = subset_or_filename
+
+    batch  = pqfile.iter_batches(batch_size=batch_size)
+    return Dataset(pa_Table.from_batches(batch))
+
+def item_to_images(subset, item):
+    """
+    Converts the images within an item (arrays), as retrieved from the dataset to proper PIL.Image
+
+    subset: The name of the Subset, one of "satellogic", "default", "sentinel-1"
+    item: The item as retrieved from the subset
+
+    returns the item, with arrays converted to PIL.Image
+    """
     metadata = item["metadata"]
     if type(metadata) == str:
         metadata = json.loads(metadata)
@@ -33,16 +87,24 @@ def item_to_images(config, item):
     }
     item["metadata"] = metadata
     
-    if config == "satellogic":
-        item["rgb"] = [
-            Image.fromarray(image.transpose(1,2,0))
-                for image in item["rgb"]
-        ]
+    if subset == "satellogic":
+        # item["rgb"] = [
+        #     Image.fromarray(np.average(image.transpose(1,2,0), 2).astype("uint8"))
+        #         for image in item["rgb"]
+        # ]
+        rgbs = []
+        for rgb in item["rgb"]:
+            rgbs.append(Image.fromarray(rgb.transpose(1,2,0)))
+            # rgbs.append(Image.fromarray(rgb[0,:,:]))      # Red
+            # rgbs.append(Image.fromarray(rgb[1,:,:]))      # Green
+            # rgbs.append(Image.fromarray(rgb[2,:,:]))      # Blue
+        item["rgb"] = rgbs
         item["1m"] = [
             Image.fromarray(image[0,:,:])
                 for image in item["1m"]
         ]
-    elif config == "sentinel_1":
+        count = len(item["1m"])
+    elif subset == "sentinel_1":
         # Mapping of V and H to RGB. May not be correct
         # https://gis.stackexchange.com/questions/400726/creating-composite-rgb-images-from-sentinel-1-channels
         i10m = item["10m"]
@@ -59,7 +121,8 @@ def item_to_images(config, item):
             Image.fromarray(image.transpose(1,2,0))
                 for image in i10m
         ]
-    elif config == "default":
+        count = len(item["10m"])
+    elif subset == "neon":
         item["rgb"] = [
             Image.fromarray(image.transpose(1,2,0))
                 for image in item["rgb"]
@@ -80,5 +143,8 @@ def item_to_images(config, item):
                 ,2).astype("uint8"))
                     for image in item["1m"]
         ]
+        count = len(item["rgb"])
+    
+    item["metadata"]["count"] = count
     return item