app.py

app.py

@@ -74,31 +74,31 @@ emb_filename = 'unsplash-25k-photos-embeddings.pkl'
 with open(emb_filename, 'rb') as emb:
         img_names, img_emb = pickle.load(emb)
  
-def display_matches(similarity):
+def display_matches(similarity, topk):
     best_matched_images = []
-    top_k_indices = torch.topk(similarity, 4, 0).indices
+    top_k_indices = torch.topk(similarity, topk, 0).indices
     for matched_image in top_k_indices:
         img = Image.open(IMAGES_DIR / img_names[matched_image])
         best_matched_images.append(img)
     return best_matched_images
   
-def image_search(search_text, search_image, option):
+def image_search(Option, topk, search_text, search_image):
   
   # Input Text Query
   #search_query = "The feeling when your program finally works"  
-  if option == "Text-To-Image" :
+  if Option == "Text-To-Image" :
     # Extracting text features embeddings
     #text_features = encode_search_query(search_text, model, device)
      text_emb = model.encode([search_text], convert_to_tensor=True)
      similarity = util.cos_sim(img_emb, text_emb)
-     return display_matches(similarity)
+     return display_matches(similarity, topk)
     # Find the matched Images
     #matched_images = find_matches(text_features, photo_features, photo_ids, 4)
      #matched_results = util.semantic_search(text_emb, img_emb, top_k=4)[0]
     
     # top 4 highest ranked images
      #return display_matches(matched_results)
-  elif option == "Image-To-Image":
+  elif Option == "Image-To-Image":
       # Input Image for Search
       #search_image = Image.fromarray(search_image.astype('uint8'), 'RGB')
       
@@ -119,13 +119,33 @@ def image_search(search_text, search_image, option):
       
       image_emb = model.encode([Image.fromarray(search_image)], convert_to_tensor=True)
       similarity = util.cos_sim(img_emb, image_emb)
-      return display_matches(similarity)
+      return display_matches(similarity, topk)
   
-gr.Interface(fn=image_search, 
-            inputs=[gr.inputs.Textbox(lines=7, label="Input Text"),
-                    gr.inputs.Image(optional=True),
-                    gr.inputs.Dropdown(["Text-To-Image", "Image-To-Image"])
+#gr.Interface(fn=image_search, 
+#            inputs=[gr.inputs.Textbox(lines=7, label="Input Text"),
+#                    gr.inputs.Image(optional=True),
+#                    gr.inputs.Dropdown(["Text-To-Image", "Image-To-Image"])
+#                    ], 
+#            outputs=gr.outputs.Carousel([gr.outputs.Image(type="pil"), gr.outputs.Image(type="pil"), gr.outputs.Image(type="pil"), gr.outputs.Image(type="pil")]),
+#            enable_queue=True
+#             ).launch(debug=True,share=True)
+gr.Interface(fn=image_search, title="Search Image",
+             description="Enter the text or image to search the other most relevant images...",
+             article=""" 
+                     Instructions:- 
+                      1. Select the option - `Text to Image`  OR `Image To Image`.
+                      2. Then accordingly enter the text or image.
+                      3. Just on entering the text or image , you will get the output image on right side
+
+                      Note: on entering the text, it may first show a different/unexpected image but then after a sec. it will show the correct image.
+                     """,
+             theme="huggingface",
+            inputs=[gr.inputs.Dropdown(["Text-To-Image", "Image-To-Image"]),
+                    gr.inputs.Dropdown(["1", "2", "3", "4", "5", "6", "7", "8", "9", "10"], type="index", default="1", label="Select Top K Images"),
+                    gr.inputs.Textbox(lines=3, label="Input Text", placeholder="Enter the text..."),
+                    gr.inputs.Image(type="pil", optional=True)
                     ], 
-            outputs=gr.outputs.Carousel([gr.outputs.Image(type="pil"), gr.outputs.Image(type="pil"), gr.outputs.Image(type="pil"), gr.outputs.Image(type="pil")]),
-            enable_queue=True
+            outputs=gr.outputs.Carousel([gr.outputs.Image(type="pil")]),
+            enable_queue=True,
+            live=True
              ).launch(debug=True,share=True)