App.py

app.py

@@ -31,7 +31,7 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 # taking features vectors
 #photo_features = np.load("./features.npy")
 
-IMAGES_DIR = './photos'
+IMAGES_DIR = Path("./photos/")
 #def show_output_image(matched_images) :
    #image=[]
    #for photo_id in matched_images:
@@ -81,19 +81,19 @@ def display_matches(indices):
 def image_search(search_text, search_image, option):
   
   # Input Text Query
-  #search_query = "The feeling when your program finally works"
-  
+  #search_query = "The feeling when your program finally works"  
   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 [Image.open(img_folder / img_names[top_k_best_image]) for top_k_best_image in torch.topk(similarity, 2, 0).indices]
     # 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]
+     #matched_results = util.semantic_search(text_emb, img_emb, top_k=4)[0]
     
     # top 4 highest ranked images
-     return display_matches(matched_results)
+     #return display_matches(matched_results)
   elif option == "Image-To-Image":
       # Input Image for Search
       #search_image = Image.fromarray(search_image.astype('uint8'), 'RGB')
@@ -107,13 +107,14 @@ def image_search(search_text, search_image, option):
       #matched_images = find_matches(image_feature, photo_ids, 4)
       
       #image_emb = model.encode(Image.open(search_image), convert_to_tensor=True)
-      image_emb = model.encode(Image.open(search_image))
+      #image_emb = model.encode(Image.open(search_image))
       # Find the matched Images
       #matched_images = find_matches(text_features, photo_features, photo_ids, 4)
       #similarity = util.cos_sim(image_emb, img_emb)
-      matched_results = util.semantic_search(image_emb, img_emb, 4)[0]
-      
-      return display_matches(matched_results)
+      #matched_results = util.semantic_search(image_emb, img_emb, 4)[0]
+      emb = model.encode([Image.fromarray(image)], convert_to_tensor=True)
+      similarity = util.cos_sim(img_emb, emb)
+      return [Image.open(IMAGES_DIR / img_names[top_k_best_image]) for top_k_best_image in torch.topk(similarity, 2, 0).indices]
   
 gr.Interface(fn=image_search, 
             inputs=[gr.inputs.Textbox(lines=7, label="Input Text"),