add nearest neighbor calculations

.gitignore

@@ -1,2 +1,3 @@
 .venv/
 __pycache__/
+.gradio/

app.py

@@ -14,6 +14,7 @@ from torchvision import transforms
 
 from templates import openai_imagenet_template
 from components.query import  get_sample
+from components.query_neighbor import QueryNeighbor
 
 log_format = "[%(asctime)s] [%(levelname)s] [%(name)s] %(message)s"
 logging.basicConfig(level=logging.INFO, format=log_format)
@@ -90,6 +91,10 @@ zero_shot_examples = [
     ],
 ]
 
+VECTOR_DB_PATH = "/Users/sreejithnoopur/codebase/bioclip-vector-db/vector_db"
+query_neighbor = QueryNeighbor(vector_db = VECTOR_DB_PATH,
+                               dataset_name = "BIRD")
+
 
 def indexed(lst, indices):
     return [lst[i] for i in indices]
@@ -146,6 +151,10 @@ def open_domain_classification(img, rank: int, return_all=False):
     logits = (model.logit_scale.exp() * img_features @ txt_emb).squeeze()
     probs = F.softmax(logits, dim=0)
 
+    neighbor = str(query_neighbor.get_nearest_neighbor(img_features))
+    neighbor_image = query_neighbor.get_image(neighbor)
+    logger.info(f"Nearest neighbor: {neighbor}")
+
     if rank + 1 == len(ranks):
         topk = probs.topk(k)
         prediction_dict = {
@@ -154,9 +163,9 @@ def open_domain_classification(img, rank: int, return_all=False):
         logger.info(f"Top K predictions: {prediction_dict}")
         top_prediction_name = format_name(*txt_names[topk.indices[0]]).split("(")[0]
         logger.info(f"Top prediction name: {top_prediction_name}")
-        sample_img, taxon_url = get_sample(metadata_df, top_prediction_name, rank)
+        _, taxon_url = get_sample(metadata_df, top_prediction_name, rank)
         if return_all:
-            return prediction_dict, sample_img, taxon_url
+            return prediction_dict, neighbor_image, taxon_url
         return prediction_dict
 
     output = collections.defaultdict(float)

components/query_neighbor.py

@@ -0,0 +1,33 @@
+import io
+import chromadb
+import boto3
+import requests
+
+from PIL import Image
+
+S3_BUCKET = "tol-bird-dataset-test"
+
+class QueryNeighbor:
+    def __init__(self, vector_db: str, dataset_name: str):
+        self._client = chromadb.PersistentClient(path=vector_db)
+        self._collection = self._client.get_collection(
+            name=dataset_name
+        )
+        self._s3_client = boto3.client("s3")
+
+
+    def get_nearest_neighbor(self, img_features) -> int:
+        ''' Returns the nearest neighbors for the given image features. '''
+        neighbors = self._collection.query(query_embeddings=[img_features[0].tolist()], 
+                     n_results = 2)
+        return neighbors["ids"][0][0]
+
+    def get_image(self, image_key: str):
+        ''' Returns the image for the given key. '''
+        img_src = self._s3_client.generate_presigned_url('get_object',
+                                                   Params={'Bucket': S3_BUCKET,
+                                                           'Key': image_key}
+                                                   )
+        img_resp = requests.get(img_src)
+        img = Image.open(io.BytesIO(img_resp.content))
+        return img