Big refactoring and extensive use of numpy

app.py

@@ -1,6 +1,7 @@
 from datasets import load_dataset, get_dataset_config_names
 from functools import partial
 from pandas import DataFrame
+from PIL import Image
 import gradio as gr 
 import numpy as np
 import tqdm
@@ -59,10 +60,74 @@ 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):
-    items = []
-    metadatas = []
+    config = state["config"]
 
+    images = []
+    metadatas = []
+    
     for i in tqdm.trange(batch_size, desc=f"Getting images"):
         if DEBUG:
             image = np.random.randint(0,255,(384,384,3))
@@ -73,40 +138,20 @@ def get_images(batch_size, state):
             except StopIteration:
                 break
             metadata = item["metadata"]
-            if state["config"] == "satellogic":
-                # image = (np.asarray(item["1m"])).astype("uint8")
-                # items.append(image[0,0,:,:])
-                image = np.asarray(item["rgb"][0]).astype(np.uint8)
-                items.append(image.transpose(1,2,0))
-
-            if state["config"] == "sentinel_1":
-                metadata = json.loads(metadata)
-                data = np.asarray(item["10m"])
-                for i in range(data.shape[0]):
-                    # 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
-                    image = np.zeros((3,384,384), "uint8")
-                    image[0] = data[i][0]
-                    image[1] = data[i][1]
-                    image[2] = (image[0]/(image[1]+0.1))*256
-                    items.append(image.transpose(1,2,0))
-
-            if state["config"] == "default":
-                dataRGB = np.asarray(item["rgb"]).astype("uint8")
-                dataCHM = np.asarray(item["chm"]).astype("uint8")
-                data1m  = np.asarray(item["1m"]).astype("uint8")
-                for i in range(dataRGB.shape[0]):
-                    image = dataRGB[i,:,:,:]
-                    items.append(image.transpose(1,2,0))
-
-                    image = dataCHM[i,0,:,:]
-                    items.append(image)
-
-                    image = data1m[i,0,:,:]
-                    items.append(image)
-            metadatas.append(metadata)
-
-    return items, DataFrame(metadatas)
+            item = item_to_images(config, item)
+
+            if  config == "satellogic":
+                images.extend(item["rgb"])
+                images.extend(item["1m"])
+            if  config == "sentinel_1":
+                images.extend(item["10m"])
+            if  config == "default":
+                images.extend(item["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)
@@ -114,53 +159,54 @@ def update_shape(rows, columns):
 def new_state():
     return gr.State({})
 
-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")
-    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,
-        interactive=False,
-        columns=5, rows=2, render=False)
-
-    with gr.Row():
-        dataset = gr.Textbox(label="Dataset", value=DATASET, interactive=False)
-        config = gr.Dropdown(choices=get_dataset_config_names(DATASET), 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")
-
-        gr.Button("Load (minutes)").click(
-            open_dataset,
-            inputs=[dataset, config, split, batch_size, state, initial_shard],
-            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, config, split, batch_size, state, shard],
-            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()
-    
-    table.render()
-
-demo.launch(show_api=False)
+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")
+        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,
+            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", )
+            split = gr.Textbox(label="Split", value="train")
+            initial_shard = gr.Number(label = "Initial shard", value=0, info="-1 for whole dataset")
+
+            gr.Button("Load (minutes)").click(
+                open_dataset,
+                inputs=[dataset, config, split, batch_size, state, initial_shard],
+                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, config, split, batch_size, state, shard],
+                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()
+        
+        table.render()
+
+    demo.launch(show_api=False)