app.py

app.py

@@ -32,9 +32,6 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 #photo_features = np.load("./features.npy")
 
 IMAGES_DIR = './photos/'
-
-
-
 #def show_output_image(matched_images) :
    #image=[]
    #for photo_id in matched_images:
@@ -77,12 +74,8 @@ emb_filename = 'unsplash-25k-photos-embeddings.pkl'
 with open(emb_filename, 'rb') as fIn:
         img_names, img_emb = pickle.load(fIn)
  
-def display_matches(similarity):
-    best_matched_images = []
-    indices = torch.topk(input=similarity, k=4, dim=0)
-    for best_img_idx in indices:
-         img = Image.open(IMAGES_DIR / img_names[best_img_idx])
-         best_matched_images.append(img)
+def display_matches(indices):
+    best_matched_images = [Image.open(os.path.join("photos/", img_names[best_img['corpus_id']])) for best_img in indices]
     return best_matched_images
   
 def image_search(search_text, search_image, option):
@@ -93,14 +86,14 @@ def image_search(search_text, search_image, option):
   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)
+     text_emb = model.encode([search_text], convert_to_tensor=True)
   
     # Find the matched Images
     #matched_images = find_matches(text_features, photo_features, photo_ids, 4)
-    similarity = util.cos_sim(text_emb, img_emb)
+     matched_results = util.semantic_search(text_emb, img_emb, 4)[0]
     
     # top 4 highest ranked images
-    return display_matches(similarity)
+     return display_matches(matched_results)
   elif option == "Image-To-Image":
       # Input Image for Search
       #search_image = Image.fromarray(search_image.astype('uint8'), 'RGB')
@@ -117,9 +110,10 @@ def image_search(search_text, search_image, option):
   
       # Find the matched Images
       #matched_images = find_matches(text_features, photo_features, photo_ids, 4)
-      similarity = util.cos_sim(image_emb, img_emb)
+      #similarity = util.cos_sim(image_emb, img_emb)
+      matched_results = util.semantic_search(image_emb, img_emb, 4)[0]
       
-      return display_matches(similarity)
+      return display_matches(matched_results)
   
 gr.Interface(fn=image_search, 
             inputs=[gr.inputs.Textbox(lines=7, label="Input Text"),