fieldvector/app.py at master · bshepp/fieldvector · GitHub

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
#!/usr/bin/env python3
"""FieldVector Gradio Web UI"""

import gradio as gr
from src.ui.state import AppState
from src.ui.components import create_field_image, create_hillshade_image, create_results_gallery
from src.config import PRESETS

state = AppState()


def load_handler(path: str, render_mode: str = 'auto'):
    if not path:
        return None, "Enter a path"
    msg = state.load_field_file(path)
    if state.field_data is not None:
        img = create_field_image(state.field_data, render_mode=render_mode)
        return img, msg
    return None, msg


def build_handler(preset: str, patch_size: int, progress=gr.Progress()):
    def cb(p, msg):
        progress(p, desc=msg)
    return state.build_new_index(preset, patch_size, 0, cb)


def query_handler(patch_id: int, k: int):
    if state.index is None:
        return None, "Load or build index first"
    msg = state.query_by_id(patch_id, k)
    if state.last_results and state.field_data is not None:
        gallery = create_results_gallery(
            state.field_data, state.last_results, state.patch_size,
            state.last_query_id, state.index.metadata.get(state.last_query_id)
        )
        return gallery, msg
    return None, msg


def create_ui():
    with gr.Blocks(title="FieldVector", theme=gr.themes.Soft()) as app:
        gr.Markdown("# 📊 FieldVector - 2D Scalar Field Similarity Search")

        with gr.Row():
            with gr.Column(scale=2):
                image_display = gr.Image(label="Field Visualization", height=500)
            with gr.Column(scale=1):
                path_input = gr.Textbox(label="Field Path")
                load_btn = gr.Button("Load")
                render_mode = gr.Dropdown(
                    ['auto', 'heatmap', 'diverging', 'hillshade', 'equalized'],
                    value='auto', label="Render Mode"
                )
                preset = gr.Dropdown(list(PRESETS.keys()), value="default", label="Preset")
                patch_size = gr.Slider(32, 256, value=64, step=32, label="Patch Size")
                build_btn = gr.Button("Build Index")
                gr.Markdown("---")
                query_id = gr.Number(label="Patch ID", value=0)
                k_results = gr.Slider(1, 50, value=8, label="Results")
                query_btn = gr.Button("Query")
                status = gr.Textbox(label="Status", interactive=False)

        gallery = gr.Gallery(label="Similar Patches", columns=4, height=300)

        load_btn.click(load_handler, [path_input, render_mode], [image_display, status])
        build_btn.click(build_handler, [preset, patch_size], [status])
        query_btn.click(query_handler, [query_id, k_results], [gallery, status])

    return app


if __name__ == '__main__':
    create_ui().launch()