Spaces:
Paused
Paused
| 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") |