fixed image similarity search

app.py

@@ -1,10 +1,11 @@
-import pickle
-import numpy as np
 import gradio as gr
+import numpy as np
+import pickle
 import clip
 import torch
 
 
+
 data = pickle.load(open("embeddings_nothumb.pkl", "rb"))
 embeddings = data["embeddings"]
 image_features = torch.Tensor(embeddings)
@@ -18,7 +19,6 @@ filenames = data["filenames"]
 base_url = "https://static.renyi.hu/ai-shared/daniel/sameenergy/index/"
 urls = [base_url + filename for filename in filenames]
 
-
 model, preprocess = clip.load('RN50')
 
 
@@ -38,20 +38,18 @@ def similarities(text_features, topk=20):
     return values, indices
 
 
-def image_retrieval(text):
+def image_retrieval_from_text(text):
     values, indices = similarities(embed_text(text), topk=20)
     top_urls = np.array(urls)[indices]
     return top_urls.tolist(), indices.numpy().tolist()
 
 
-def on_select(evt):
-    print("event:", evt)
-    return str(evt)
-    return f"You selected {evt.value} at {evt.index} from {evt.target}"
-
-
-def empty_gallery():
-    return [], []
+def image_retrieval_from_image(state, selected_locally):
+    selected = state[int(selected_locally)]
+    image_vector = image_features[selected][None, :]
+    values, indices = similarities(image_vector, topk=20)
+    top_urls = np.array(urls)[indices]
+    return top_urls.tolist(), indices.numpy().tolist()
 
 
 with gr.Blocks(css="footer {visibility: hidden}") as demo:
@@ -64,21 +62,25 @@ with gr.Blocks(css="footer {visibility: hidden}") as demo:
             max_lines=1,
             placeholder="Enter your prompt",
         ).style(container=False)
-        btn = gr.Button("Search").style(full_width=False)
-
+        text_query_button = gr.Button("Search").style(full_width=False)
 
     gallery = gr.Gallery(label="Images", show_label=False, elem_id="gallery"
         ).style(columns=4, container=False)
 
-    demo.load(empty_gallery, None, [gallery, state])
+    # demo.load(empty_gallery, None, [gallery, state])
+
+    with gr.Row():
+        selected = gr.Number(placeholder=0, show_label=False)
+        image_query_button = gr.Button("Show similar")
 
-    selected = gr.Textbox(placeholder="Selected", show_label=False)
+    text_query_button.click(image_retrieval_from_text, [text], [gallery, state])
+    image_query_button.click(image_retrieval_from_image, [state, selected], [gallery, state])
 
-    btn.click(image_retrieval, text, [gallery, state])
+    def get_select_index(evt: gr.SelectData):
+        return evt.index
 
-    # does not work, function is called with None instead of event:
-    gallery.select(on_select, None, selected)
+    gallery.select(get_select_index, None, selected)
 
 
 if __name__ == "__main__":
-    demo.launch(height=2000)
+    demo.launch()