File size: 3,240 Bytes
465c443
63711ea
 
 
c501c19
 
465c443
63711ea
 
 
7a6ee1c
7c425bb
63711ea
 
 
 
 
 
 
 
 
 
 
465c443
63711ea
465c443
63711ea
 
 
 
 
 
c501c19
d3015aa
7c425bb
c501c19
 
63711ea
c501c19
63711ea
c501c19
 
99a41a1
c501c19
99a41a1
63711ea
7a6ee1c
c501c19
63711ea
 
c501c19
63711ea
 
c501c19
 
 
30231d8
99a41a1
c501c19
 
 
 
99a41a1
63711ea
 
 
c501c19
 
7a6ee1c
99a41a1
 
63711ea
99a41a1
63711ea
99a41a1
 
 
 
 
 
 
63711ea
99a41a1
 
 
4561290
99a41a1
 
 
 
 
 
 
 
 
 
63711ea
 
 
 
 
 
 
 
30231d8
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
import gradio as gr
import earthview as ev
import os
import json
import utils
from pandas import DataFrame

# --- Configuration ---
DATASET_SUBSET = "satellogic"
LABELED_DATA_FILE = "labeled_data.json"
SAMPLE_SEED = 10  # The seed to use when sampling the dataset for the demo.

# --- Load Dataset ---
dataset = ev.load_dataset(DATASET_SUBSET, shards=[SAMPLE_SEED])
data_iter = iter(dataset)

# --- Load Labeled Data (if it exists) ---
def load_labeled_data():
    if os.path.exists(LABELED_DATA_FILE):
        with open(LABELED_DATA_FILE, "r") as f:
            return json.load(f)
    else:
        return []

labeled_data = load_labeled_data()

# --- Get Next Sample for Labeling ---
def get_next_sample():
    global data_iter
    try:
        sample = next(data_iter)
        sample = ev.item_to_images(DATASET_SUBSET, sample)
        return sample
    except StopIteration:
        print("No more samples in the dataset.")
        return None

# --- Save Labeled Data ---
def save_labeled_data(label, state):
    global labeled_data
    sample = state["sample"]
    if sample is None:
        return "No image to label", DataFrame()

    image_bytes = sample["rgb"][0].convert("RGB").tobytes()  # Convert PIL Image to bytes
    labeled_data.append({
        "image": image_bytes,
        "metadata": sample["metadata"],
        "label": label
    })

    with open(LABELED_DATA_FILE, "w") as f:
        json.dump(labeled_data, f)

    new_sample = get_next_sample()
    if new_sample is None:
        state["sample"] = None
        return "Dataset exhausted.", DataFrame()

    state["sample"] = new_sample
    new_metadata = new_sample["metadata"]
    new_metadata["map"] = f'<a href="{utils.get_google_map_link(new_sample, DATASET_SUBSET)}" target="_blank">🧭</a>'
    return "", DataFrame([new_metadata])

# --- Gradio Interface ---
def labeling_ui():
    state = gr.State({"sample": None, "subset": DATASET_SUBSET})

    with gr.Row():
        cool_button = gr.Button("Cool")
        not_cool_button = gr.Button("Not Cool")

    table = gr.DataFrame(datatype="html")

    def initialize_labeling_ui():
        sample = get_next_sample()
        if sample is None:
            return {"sample": None, "subset": DATASET_SUBSET}, DataFrame()
        metadata = sample["metadata"]
        metadata["map"] = f'<a href="{utils.get_google_map_link(sample, DATASET_SUBSET)}" target="_blank">🧭</a>'
        return {"sample": sample, "subset": DATASET_SUBSET}, DataFrame([metadata])

    initial_state, initial_metadata = initialize_labeling_ui()
    table.value = initial_metadata
    state.value = initial_state

    cool_button.click(
        fn=lambda label, state: save_labeled_data(label, state),
        inputs=[gr.Textbox(visible=False, value="cool"), state],
        outputs=[gr.Textbox(label="Debug"), table]
    )
    not_cool_button.click(
        fn=lambda label, state: save_labeled_data(label, state),
        inputs=[gr.Textbox(visible=False, value="not cool"), state],
        outputs=[gr.Textbox(label="Debug"), table]
    )

# --- Main Interface ---
with gr.Blocks() as demo:
    gr.Markdown("# TerraNomaly")
    with gr.Tabs():
        with gr.TabItem("Labeling"):
            labeling_ui()

demo.launch(debug=True)