s1ngledoge's picture
upd
1c8d1ba
Raw
History Blame Contribute Delete
8.68 kB
from __future__ import annotations
from typing import Any
import gradio as gr
from src.ui.callbacks import AppCallbacks, current_categories
from src.ui.constants import (
ALL_CATEGORIES_LABEL,
APP_NAME,
APP_SUBTITLE,
DEFAULT_MIN_SIMILARITY,
DEFAULT_TOP_K,
EXAMPLE_QUERIES,
IMAGE_QUERY_EMPTY_PREVIEW,
MAX_TOP_K,
NO_FAVORITES_STATUS,
TEXT_QUERY_EMPTY_PREVIEW,
gradio_launch_options,
)
from src.ui.formatters import initial_overview
def create_demo(callbacks: AppCallbacks) -> tuple[gr.Blocks, dict[str, Any]]:
blocks_kwargs, launch_kwargs = gradio_launch_options()
with gr.Blocks(**blocks_kwargs) as demo:
current_results_state = gr.State([])
selected_result_state = gr.State(None)
favorites_state = gr.State([])
query_mode_state = gr.State("text")
with gr.Column(elem_id="app-shell"):
gr.Markdown(
f"# {APP_NAME}\n\n{APP_SUBTITLE}",
elem_id="hero-section",
)
with gr.Row(elem_id="main-dashboard", equal_height=False):
with gr.Column(scale=1, min_width=300, elem_id="left-sidebar"):
with gr.Column(elem_id="search-card", elem_classes=["search-card"]):
with gr.Tabs(elem_id="query-tabs"):
with gr.Tab("Text Search") as text_tab:
text_query = gr.Textbox(
label="Image Description",
placeholder="For example: red apple, milk bottle, green vegetables",
lines=2,
elem_id="text-query",
)
gr.Examples(
examples=EXAMPLE_QUERIES,
inputs=text_query,
label="Example Queries",
elem_id="example-chips",
)
text_preview = gr.Markdown(
TEXT_QUERY_EMPTY_PREVIEW,
elem_id="text-preview",
)
with gr.Tab("Image Search") as image_tab:
image_query = gr.Image(
label="Upload Query Image",
type="pil",
placeholder="Drag an image here, or click to upload",
sources=["upload", "clipboard"],
height=210,
elem_id="image-query",
)
image_preview = gr.Markdown(
IMAGE_QUERY_EMPTY_PREVIEW,
elem_id="image-preview",
)
with gr.Column(elem_id="parameter-card", elem_classes=["parameter-card"]):
gr.Markdown("### Search Parameters", elem_id="parameter-title")
top_k = gr.Slider(
minimum=1,
maximum=MAX_TOP_K,
value=DEFAULT_TOP_K,
step=1,
label="Number of Results",
elem_id="top-k",
)
category_filter = gr.Dropdown(
choices=current_categories(),
value=ALL_CATEGORIES_LABEL,
label="Category Filter",
allow_custom_value=False,
elem_id="category-filter",
)
min_similarity = gr.Slider(
minimum=0,
maximum=1,
value=DEFAULT_MIN_SIMILARITY,
step=0.01,
label="Minimum Similarity",
elem_id="min-similarity",
)
with gr.Row(elem_id="action-row"):
search_button = gr.Button(
"Search",
variant="primary",
elem_id="search-button",
elem_classes=["primary-button"],
)
clear_button = gr.Button(
"Clear Query",
elem_id="clear-button",
elem_classes=["secondary-button"],
)
with gr.Column(scale=3, min_width=360, elem_id="right-workspace"):
with gr.Column(elem_id="result-card", elem_classes=["result-card"]):
result_overview = gr.Markdown(
initial_overview(),
elem_id="result-overview",
)
result_gallery = gr.Gallery(
label="Search Results",
columns=4,
height=390,
object_fit="cover",
show_label=True,
elem_id="result-gallery",
elem_classes=["result-gallery"],
)
with gr.Column(elem_id="favorite-card", elem_classes=["favorite-card"]):
with gr.Row(elem_id="favorite-header"):
gr.Markdown("### My Favorites", elem_id="favorite-title")
favorite_status = gr.Markdown(NO_FAVORITES_STATUS, elem_id="favorite-status")
with gr.Row(elem_id="favorite-actions"):
add_favorite_button = gr.Button(
"Add to Favorites",
elem_id="add-favorite-button",
elem_classes=["secondary-button"],
)
export_button = gr.DownloadButton(
label="Export Favorites as ZIP",
value=None,
elem_id="export-button",
elem_classes=["primary-button"],
)
favorite_gallery = gr.Gallery(
label="My Favorites",
columns=5,
height=180,
object_fit="cover",
show_label=True,
elem_id="favorite-gallery",
elem_classes=["favorite-gallery"],
)
text_tab.select(callbacks.select_text_mode, outputs=query_mode_state)
image_tab.select(callbacks.select_image_mode, outputs=query_mode_state)
text_query.change(callbacks.preview_text_query, inputs=text_query, outputs=text_preview)
image_query.change(callbacks.preview_image_query, inputs=image_query, outputs=image_preview)
search_button.click(
callbacks.run_search,
inputs=[query_mode_state, text_query, image_query, top_k, category_filter, min_similarity],
outputs=[result_overview, result_gallery, selected_result_state, current_results_state],
)
clear_button.click(
callbacks.clear_search,
inputs=[top_k, category_filter, min_similarity],
outputs=[
text_query,
image_query,
text_preview,
image_preview,
result_overview,
result_gallery,
selected_result_state,
current_results_state,
],
)
result_gallery.select(
callbacks.select_result_from_gallery,
inputs=current_results_state,
outputs=selected_result_state,
show_progress="hidden",
)
add_favorite_button.click(
callbacks.add_selected_to_favorites,
inputs=[selected_result_state, current_results_state, favorites_state],
outputs=[favorites_state, favorite_gallery, favorite_status, export_button],
)
return demo, launch_kwargs