File size: 5,646 Bytes
9efa0c5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import time
from pathlib import Path

import gradio as gr

from qa import ask_about_image
from search import search, filenames
from ui.theme import theme, CSS, TOGGLE_DARK


# ============================================================
# Config
# ============================================================

IMAGE_DIR = Path("data/Flicker8k_Dataset")
EXAMPLES = [
    "a dog running on the beach",
    "children playing soccer",
    "a man climbing a rock",
    "a person riding a bike in the snow",
]

SCROLL_TO_ASK = """
() => document.getElementById('ask-panel')
        ?.scrollIntoView({ behavior: 'smooth', block: 'center' })
"""

ANSWER_PLACEHOLDER = "Ask a question about this image."

# ============================================================
# Handlers
# ============================================================

def run_search(query: str, k: int):
    if not query.strip():
        return (gr.update(value=[], visible=False), "", gr.update(visible=True),
                [], gr.update(visible=False))

    start = time.perf_counter()
    results = search(query, k=int(k))
    elapsed = time.perf_counter() - start

    paths, items = [], []
    for rank, (name, score) in enumerate(results, start=1):
        path = str(IMAGE_DIR / name)
        paths.append(path)
        items.append((path, f"#{rank} 路 {score:.3f}"))

    status_text = f"{len(results)} of {len(filenames):,} 路 {elapsed:.2f} s"
    return (gr.update(value=items, visible=True), status_text, gr.update(visible=False),
            paths, gr.update(visible=False))


def on_select(evt: gr.SelectData, paths: list[str]):
    path = paths[evt.index]
    info = f"Result #{evt.index + 1} 路 `{Path(path).name}`"
    return gr.update(visible=True), path, path, info, "", ANSWER_PLACEHOLDER

def on_ask(image_path: str | None, question: str):
    if not image_path:
        return "Click an image in the results first."
    return ask_about_image(image_path, question)

# ============================================================
# Layout
# ============================================================

with gr.Blocks(title="Image Search") as demo:
    result_paths = gr.State([])
    selected_path = gr.State(None)

    # --- Header ---
    with gr.Row(elem_id="header"):
        gr.Markdown(
            "# Image Search\nDescribe a scene. We rank the photos by similarity.",
            elem_id="title",
        )
        theme_btn = gr.Button("Light / Dark", size="sm", variant="secondary", elem_id="theme-btn")

    # --- Search ---
    with gr.Column(elem_classes="card"):
        with gr.Row(elem_id="search-row"):
            query = gr.Textbox(label="Query", placeholder="a dog running on the beach", scale=5)
            k = gr.Slider(1, 20, value=8, step=1, label="Top-k", scale=2)
            button = gr.Button("Search", variant="primary", scale=1, min_width=120, elem_id="search-btn")
        with gr.Row(elem_id="chips"):
            gr.Markdown("Try:", elem_id="try-label")
            chips = [gr.Button(ex, size="sm", variant="secondary", elem_classes="chip") for ex in EXAMPLES]

    # --- Results ---
    with gr.Row(elem_id="results-head"):
        gr.Markdown("## Results")
        status = gr.Markdown("", elem_id="status")

    empty = gr.Markdown("Your search results will appear here", elem_id="empty-state")

    gallery = gr.Gallery(
        show_label=False,
        container=False,
        columns=4,
        height="auto",
        object_fit="cover",
        allow_preview=False,
        elem_id="gallery",
        visible=False,
    )

    # --- Ask about this image ---
    with gr.Column(elem_classes="card", elem_id="ask-panel", visible=False) as ask_panel:
        with gr.Row(elem_id="ask-row"):
            selected_img = gr.Image(
                show_label=False, interactive=False, height=220,
                scale=1, min_width=220, elem_id="ask-image",
            )
            with gr.Column(scale=3):
                gr.Markdown("### Ask about this image")
                selected_info = gr.Markdown("", elem_id="ask-info")
                with gr.Row(elem_id="ask-input-row"):
                    question = gr.Textbox(
                        show_label=False, placeholder="e.g. What color is the dog?",
                        scale=5,
                    )
                    ask_btn = gr.Button("Ask", variant="primary", scale=1, min_width=100)
                answer = gr.Markdown(ANSWER_PLACEHOLDER, elem_id="ask-answer")


    # ========================================================
    # Events
    # ========================================================

    search_outputs = [gallery, status, empty, result_paths, ask_panel]

    theme_btn.click(None, js=TOGGLE_DARK)
    button.click(run_search, inputs=[query, k], outputs=search_outputs)
    query.submit(run_search, inputs=[query, k], outputs=search_outputs)
    for chip, ex in zip(chips, EXAMPLES):
        chip.click(lambda q=ex: q, outputs=query).then(
            run_search, inputs=[query, k], outputs=search_outputs
        )

    gallery.select(
        on_select, inputs=[result_paths],
        outputs=[ask_panel, selected_img, selected_path, selected_info, question, answer],
    ).then(None, js=SCROLL_TO_ASK)

    ask_btn.click(on_ask, inputs=[selected_path, question], outputs=answer, concurrency_limit=2)
    question.submit(on_ask, inputs=[selected_path, question], outputs=answer, concurrency_limit=2)

# ============================================================
# Run
# ============================================================

if __name__ == "__main__":
    demo.launch(theme=theme, css=CSS, server_name="0.0.0.0")