|
import gradio as gr |
|
import os |
|
from qdrant_client import QdrantClient |
|
from sentence_transformers import SentenceTransformer |
|
|
|
model = SentenceTransformer("clip-ViT-B-32") |
|
|
|
qdrant_client = QdrantClient( |
|
url = os.environ['QDRANT_URL'], |
|
port= 443, |
|
api_key = os.environ['QDRANT_API_KEY'] |
|
) |
|
|
|
def search_images(modality, count, input_text, input_image): |
|
query = str(input_text) if modality=='Text' else input_image |
|
|
|
results = qdrant_client.search( |
|
collection_name = "images", |
|
query_vector = model.encode(query).tolist(), |
|
with_payload = True, |
|
limit = count |
|
) |
|
|
|
return [gr.update(value="## Results\nThe image data is limited, don't expect to find everything!")]+[gr.Image(value=result.payload['url'], visible=True) for result in results]+[gr.Image(visible=False)]*(100-count) |
|
|
|
def clear(): |
|
return [gr.update(value="")]+[gr.Image(visible=False)]*100 |
|
|
|
def input_interface(choice): |
|
if choice == "Text": |
|
return [gr.update(visible=True), gr.update(visible=False)] |
|
else: |
|
return [gr.update(visible=False), gr.update(visible=True)] |
|
|
|
with gr.Blocks() as interface: |
|
gr.Markdown("# Multi-Modal Image Search Engine\nSemantically search over 15k images using text or image inputs!") |
|
|
|
|
|
with gr.Column(variant='compact'): |
|
input_type = gr.Radio(choices=["Text", "Image"], type="value", label="Modality", value="Text") |
|
with gr.Column() as text_area: |
|
text_input = gr.Textbox(label="Text", lines=1, placeholder="Try 'Golden Retriever'") |
|
with gr.Column(visible=False) as image_uploader: |
|
image_input = gr.Image(type="pil") |
|
input_type.change(input_interface, input_type, [text_area, image_uploader]) |
|
|
|
|
|
with gr.Column(variant="panel"): |
|
count = gr.Slider(minimum=1, maximum=40, step=1, value=8, label="No. of Results") |
|
images_btn = gr.Button(value="Search Images", variant="primary") |
|
|
|
|
|
images = [] |
|
images.append(gr.Markdown()) |
|
with gr.Column() as output_images: |
|
for i in range(10): |
|
with gr.Row(): |
|
for j in range(4): |
|
images.append(gr.Image(visible=False)) |
|
images_btn.click(clear, outputs=images).then(search_images, inputs=[input_type, count, text_input, image_input], outputs=images) |
|
|
|
interface.launch() |
|
|