moved code to earthview.py

app.py

@@ -1,29 +1,15 @@
 from datasets import load_dataset, get_dataset_config_names
 from functools import partial
 from pandas import DataFrame
-from PIL import Image
+import earthview as ev
 import gradio as gr 
-import numpy as np
 import tqdm
-import json
 import os
 
-DATASET = "satellogic/EarthView"
 DEBUG = False
 
-sets = {
-    "satellogic": {
-        "shards" : 3676,
-    },
-    "sentinel_1": {
-        "shards" : 1763,
-    },
-    "neon": {
-        "config" : "default",
-        "shards" : 607,
-        "path"   : "data",
-    }
-}
+if DEBUG:
+    import numpy as np
 
 def open_dataset(dataset, set_name, split, batch_size, state, shard = -1):
     if shard == -1:
@@ -31,9 +17,9 @@ def open_dataset(dataset, set_name, split, batch_size, state, shard = -1):
         data_files = None
         shards = 100
     else:
-        config = sets[set_name].get("config", set_name)
-        shards = sets[set_name]["shards"]
-        path   = sets[set_name].get("path", set_name)
+        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"]}
 
     if DEBUG:
@@ -60,68 +46,6 @@ def open_dataset(dataset, set_name, split, batch_size, state, shard = -1):
         state
     )
     
-def item_to_images(config, item):
-    metadata = item["metadata"]
-    if type(metadata) == str:
-        metadata = json.loads(metadata)
-
-    item = {
-        k: np.asarray(v).astype("uint8")
-            for k,v in item.items()
-                if k != "metadata"
-    }
-    item["metadata"] = metadata
-    
-    if config == "satellogic":
-        item["rgb"] = [
-            Image.fromarray(image.transpose(1,2,0))
-                for image in item["rgb"]
-        ]
-        item["1m"] = [
-            Image.fromarray(image[0,:,:])
-                for image in item["1m"]
-        ]
-    elif config == "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"]
-        i10m = np.concatenate(
-            (   i10m,
-                np.expand_dims(
-                    i10m[:,0,:,:]/(i10m[:,1,:,:]+0.01)*256,
-                    1
-                ).astype("uint8")
-            ),
-            1
-        )
-        item["10m"] = [
-            Image.fromarray(image.transpose(1,2,0))
-                for image in i10m
-        ]
-    elif config == "default":
-        item["rgb"] = [
-            Image.fromarray(image.transpose(1,2,0))
-                for image in item["rgb"]
-        ]
-        item["chm"] = [
-            Image.fromarray(image[0])
-                for image in item["chm"]
-        ]
-
-        # The next is a very arbitrary conversion from the 369 hyperspectral data to RGB
-        # It just averages each 1/3 of the bads and assigns it to a channel
-        item["1m"] = [
-            Image.fromarray(
-                np.concatenate((
-                    np.expand_dims(np.average(image[:124],0),2),
-                    np.expand_dims(np.average(image[124:247],0),2),
-                    np.expand_dims(np.average(image[247:],0),2))
-                ,2).astype("uint8"))
-                    for image in item["1m"]
-        ]
-    return item
-
-
 def get_images(batch_size, state):
     config = state["config"]
 
@@ -138,7 +62,7 @@ def get_images(batch_size, state):
             except StopIteration:
                 break
             metadata = item["metadata"]
-            item = item_to_images(config, item)
+            item = ev.item_to_images(config, item)
 
             if  config == "satellogic":
                 images.extend(item["rgb"])
@@ -163,20 +87,20 @@ if __name__ == "__main__":
     with gr.Blocks(title="Dataset Explorer", fill_height = True) as demo:
         state = new_state()
 
-        gr.Markdown(f"# Viewer for [{DATASET}](https://huggingface.co/datasets/satellogic/EarthView) Dataset")
+        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=DATASET,
+            label=ev.DATASET,
             interactive=False,
             columns=5, rows=2, render=False)
 
         with gr.Row():
-            dataset = gr.Textbox(label="Dataset", value=DATASET, interactive=False)
-            config = gr.Dropdown(choices=sets.keys(), label="Config", value="satellogic", )
+            dataset = gr.Textbox(label="Dataset", value=ev.DATASET, interactive=False)
+            config = gr.Dropdown(choices=ev.get_sets(), label="Config", value="satellogic", )
             split = gr.Textbox(label="Split", value="train")
             initial_shard = gr.Number(label = "Initial shard", value=0, info="-1 for whole dataset")
 

earthview.py

@@ -0,0 +1,84 @@
+from PIL import Image
+import numpy as np
+import json
+
+DATASET = "satellogic/EarthView"
+
+sets = {
+    "satellogic": {
+        "shards" : 3676,
+    },
+    "sentinel_1": {
+        "shards" : 1763,
+    },
+    "neon": {
+        "config" : "default",
+        "shards" : 607,
+        "path"   : "data",
+    }
+}
+
+def get_sets():
+    return sets.keys()
+
+def item_to_images(config, item):
+    metadata = item["metadata"]
+    if type(metadata) == str:
+        metadata = json.loads(metadata)
+
+    item = {
+        k: np.asarray(v).astype("uint8")
+            for k,v in item.items()
+                if k != "metadata"
+    }
+    item["metadata"] = metadata
+    
+    if config == "satellogic":
+        item["rgb"] = [
+            Image.fromarray(image.transpose(1,2,0))
+                for image in item["rgb"]
+        ]
+        item["1m"] = [
+            Image.fromarray(image[0,:,:])
+                for image in item["1m"]
+        ]
+    elif config == "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"]
+        i10m = np.concatenate(
+            (   i10m,
+                np.expand_dims(
+                    i10m[:,0,:,:]/(i10m[:,1,:,:]+0.01)*256,
+                    1
+                ).astype("uint8")
+            ),
+            1
+        )
+        item["10m"] = [
+            Image.fromarray(image.transpose(1,2,0))
+                for image in i10m
+        ]
+    elif config == "default":
+        item["rgb"] = [
+            Image.fromarray(image.transpose(1,2,0))
+                for image in item["rgb"]
+        ]
+        item["chm"] = [
+            Image.fromarray(image[0])
+                for image in item["chm"]
+        ]
+
+        # The next is a very arbitrary conversion from the 369 hyperspectral data to RGB
+        # It just averages each 1/3 of the bads and assigns it to a channel
+        item["1m"] = [
+            Image.fromarray(
+                np.concatenate((
+                    np.expand_dims(np.average(image[:124],0),2),
+                    np.expand_dims(np.average(image[124:247],0),2),
+                    np.expand_dims(np.average(image[247:],0),2))
+                ,2).astype("uint8"))
+                    for image in item["1m"]
+        ]
+    return item
+