app.py

app.py

@@ -75,7 +75,7 @@ with open(emb_filename, 'rb') as fIn:
         img_names, img_emb = pickle.load(fIn)
  
 def display_matches(indices):
-    best_matched_images = [Image.open(os.path.join("photos/", img_names[best_img['corpus_id']])) for best_img in indices]
+    best_matched_images = [Image.open(IMAGES_DIR / img_names[top_k_best_image]) for top_k_best_image in torch.topk(similarity, 2, 0).indices]
     return best_matched_images
   
 def image_search(search_text, search_image, option):
@@ -87,7 +87,7 @@ def image_search(search_text, search_image, option):
     #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(IMAGES_DIR / img_names[top_k_best_image]) for top_k_best_image in torch.topk(similarity, 2, 0).indices]
+     return display_matches(similarity)
     # 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]
@@ -112,9 +112,9 @@ def image_search(search_text, search_image, option):
       #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]
-      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]
+      imageg_emb = model.encode([Image.fromarray(search_image)], convert_to_tensor=True)
+      similarity = util.cos_sim(img_emb, image_emb)
+      return display_matches(similarity)
   
 gr.Interface(fn=image_search, 
             inputs=[gr.inputs.Textbox(lines=7, label="Input Text"),