File size: 2,206 Bytes
2a4ddda
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import gradio as gr
from qdrant_client import QdrantClient
from sentence_transformers import SentenceTransformer

model = SentenceTransformer("clip-ViT-B-32")

qdrant_client = QdrantClient(
    url = 'url',
    port= 443,
    api_key = "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")]+[gr.update(value=result.payload['url'], visible=True) for result in results]+[gr.update(visible=False)]*(100-count)

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. The image data is limited, don't expect to find everything!")

    # Input Interface
    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])

    # Search Controls
    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")

    # Output Interface
    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(search_images, inputs=[input_type, count, text_input, image_input], outputs=images)

    interface.launch()