File size: 5,250 Bytes
6d787c4
 
 
c670c44
 
6d787c4
 
c670c44
6d787c4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
41c9d30
 
6d787c4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c670c44
4779900
6d787c4
 
 
 
 
 
 
 
 
41c9d30
6d787c4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4779900
 
6d787c4
 
4779900
6d787c4
 
 
 
 
 
 
 
 
 
 
4779900
6d787c4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3326301
6d787c4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
11296d8
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
from datasets import load_dataset, IterableDataset
from functools import partial
from pandas import DataFrame
import gradio as gr 
import numpy as np
import tqdm
import json
import os

DEBUG = False

sets = {
    "satellogic": {
        "shards" : 3676,
    },
    "sentinel_1": {
        "shards" : 1763,
    },
    "neon": {
        "config" : "default",
        "shards" : 607,
        "path"   : "data",
    }
}

def open_dataset(dataset, set_name, split, batch_size, shard = -1):
    # I should really move ds.config_name and dsi to a gr.State()
    global dsi, ds

    if 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)
        data_files = {"train":[f"{path}/{split}-{shard:05d}-of-{shards:05d}.parquet"]}

    if DEBUG:
        ds  = lambda:None
        ds.n_shards = 1234
        dsi = range(100)
    else:
        ds = load_dataset(
            dataset,
            config,
            split=split,
            cache_dir="dataset",
            data_files=data_files,
            streaming=True,
            token=os.environ.get("HF_TOKEN", None))
    
        dsi = iter(ds)

    return (
        gr.update(label=f"Shards (max {shards})", value=shard, maximum=shards),
        *get_images(batch_size)
    )
    
def get_images(batch_size):
    global dsi, ds
        
    items = []
    metadatas = []

    for i in tqdm.trange(batch_size, desc=f"Getting images"):
        if DEBUG:
            image = np.random.randint(0,255,(384,384,3))
            metadata = {"bounds":[[1,1,4,4]], }
        else:
            try:
                item = next(dsi)
            except StopIteration:
                break
            metadata = item["metadata"]
            if ds.config_name == "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 ds.config_name == "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 ds.config_name == "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)

def update_shape(rows, columns):
    return gr.update(rows=rows, columns=columns)


with gr.Blocks(title="Dataset Explorer", fill_height = True) as demo:
    gr.Markdown("# [satellogic/EarthView](https://huggingface.co/datasets/satellogic/EarthView) Dataset Viewer")
    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="satellogic/EarthView",
        interactive=False,
        columns=5, rows=2, render=False)

    with gr.Row():
        dataset = gr.Textbox(label="Dataset", value="satellogic/EarthView")
        config = gr.Dropdown(choices=["satellogic", "sentinel_1", "neon"], label="Subset", value="satellogic", )
        split = gr.Textbox(label="Split", value="train")
        initial_shard = gr.Number(label = "Initial shard", value=0)

        gr.Button("Load (minutes)").click(
            open_dataset,
            inputs=[dataset, config, split, batch_size, initial_shard],
            outputs=[shard, gallery, table])

    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, shard],
            outputs=[shard, gallery, table])

        btn = gr.Button("Get More Images", scale=0)
        btn.click(get_images, [batch_size], [gallery, table])
        btn.click()
    
    table.render()

demo.launch(show_api=False)