Update src/app.py

src/app.py

@@ -87,24 +87,46 @@ def get_neighbors(query_image, selected_descriptor, selected_distance, top_k=5):
     if "CLIP" == selected_descriptor:
         clip_model = CLIPImageEncoder()
         qi_embedding = clip_model.encode_image(query_image)
-        scores, retrieved_examples = dataset_with_embeddings.get_nearest_examples(
-            'clip_embeddings', qi_embedding, k=top_k)
+        if selected_distance == "FAISS":
+            scores, retrieved_examples = dataset_with_embeddings.get_nearest_examples(
+                'clip_embeddings', qi_np, k=top_k)
+        elif selected_distance == "Chi-squared":
+            tmp_dataset = dataset_with_embeddings.map(lambda row: {'distance': chi2_distance(histA=qi_embedding, histB=row['clip_embeddings'])})
+            retrieved_examples = tmp_dataset.sort("distance")[:5]
+        else:
+            tmp_dataset = dataset_with_embeddings.map(lambda row: {'distance': euclidean_distance(histA=qi_embedding, histB=row['clip_embeddings'])})
+            retrieved_examples = tmp_dataset.sort("distance")[:5]
         images = retrieved_examples['image']
         return images
     if "LBP" == selected_descriptor:
         lbp_model = LBPImageEncoder(8,2)
         qi_embedding = lbp_model.describe(query_image)
-        scores, retrieved_examples = dataset_with_embeddings.get_nearest_examples(
-            'lbp_embeddings', qi_embedding, k=top_k)
+        if selected_distance == "FAISS":
+            scores, retrieved_examples = dataset_with_embeddings.get_nearest_examples(
+                'lbp_embeddings', qi_np, k=top_k)
+        elif selected_distance == "Chi-squared":
+            tmp_dataset = dataset_with_embeddings.map(lambda row: {'distance': chi2_distance(histA=qi_embedding, histB=row['lbp_embeddings'])})
+            retrieved_examples = tmp_dataset.sort("distance")[:5]
+        else:
+            tmp_dataset = dataset_with_embeddings.map(lambda row: {'distance': euclidean_distance(histA=qi_embedding, histB=row['lbp_embeddings'])})
+            retrieved_examples = tmp_dataset.sort("distance")[:5]
         images = retrieved_examples['image']
         return images
     if "LBPColor" == selected_descriptor:
         lbp_model = LBPImageEncoder(8,2)
         cd = ColorDescriptor((8, 12, 3))
         qi_embedding = merge_features(lbp_model.describe(query_image), cd.describe(query_image))
-        scores, retrieved_examples = dataset_with_embeddings.get_nearest_examples(
-            'lbp_color_embeddings', qi_embedding, k=top_k
-        )
+        if selected_distance == "FAISS":
+            scores, retrieved_examples = dataset_with_embeddings.get_nearest_examples(
+                'lbp_color_embeddings', qi_np, k=top_k)
+        elif selected_distance == "Chi-squared":
+            tmp_dataset = dataset_with_embeddings.map(lambda row: {'distance': chi2_distance(histA=qi_embedding, histB=row['lbp_color_embeddings'])})
+            retrieved_examples = tmp_dataset.sort("distance")[:5]
+        else:
+            tmp_dataset = dataset_with_embeddings.map(lambda row: {'distance': euclidean_distance(histA=qi_embedding, histB=row['lbp_color_embeddings'])})
+            retrieved_examples = tmp_dataset.sort("distance")[:5]
+        images = retrieved_examples['image']
+        return images
 
 
     else: