dinoV2

configs/embed_folder.yml

@@ -0,0 +1,7 @@
+command:
+  (): geoguessr_bot.commands.EmbedCommand
+  embedder:
+    (): geoguessr_bot.retriever.DinoV2Embedder
+    device: "cpu"
+  images_folder: !path "../data/samples"
+  output_path: !path "../data/samples_embedded"

geoguessr_bot/__main__.py

@@ -0,0 +1,17 @@
+import fire
+from configue import load
+
+
+class CLI:
+    """Regroup all the commands of the CLI
+    """
+
+    @staticmethod
+    def run(config_path: str) -> None:
+        config = load(config_path)
+        command = config["command"]
+        command.run()
+
+
+if __name__ == "__main__":
+    fire.Fire(CLI)

geoguessr_bot/commands/__init__.py

@@ -0,0 +1,2 @@
+from .abstract_command import AbstractCommand
+from .embed_command import EmbedCommand

geoguessr_bot/commands/abstract_command.py

@@ -0,0 +1,9 @@
+from abc import abstractmethod
+
+
+class AbstractCommand:
+    """Abstract class for all commands.
+    """
+    @abstractmethod
+    def run(self) -> None:
+        raise NotImplementedError

geoguessr_bot/commands/embed_command.py

@@ -0,0 +1,16 @@
+from dataclasses import dataclass
+
+from geoguessr_bot.commands import AbstractCommand
+from geoguessr_bot.retriever import AbstractImageEmbedder
+
+
+@dataclass
+class EmbedCommand(AbstractCommand):
+    """Embed all images in a folder and save them in a .npy file
+    """
+    embedder: AbstractImageEmbedder
+    images_folder: str
+    output_path: str
+
+    def run(self) -> None:
+        self.embedder.embed_folder(self.images_folder, self.output_path)

geoguessr_bot/retriever/__init__.py

@@ -1,3 +1,4 @@
-from .retriever import AbstractRetriever
+from .retriever import Retriever
 from .abstract_embedder import AbstractImageEmbedder
 from .dino_embedder import DinoEmbedder
+from .dino_v2_embedder import DinoV2Embedder

geoguessr_bot/retriever/abstract_embedder.py

@@ -1,6 +1,7 @@
 import os
 from PIL import Image
 import numpy as np
+from tqdm import tqdm
 
 
 class AbstractImageEmbedder:
@@ -12,13 +13,14 @@ class AbstractImageEmbedder:
         """
         raise NotImplementedError
 
-    def embed_folder(self, folder_path: str):
+    def embed_folder(self, folder_path: str, output_path: str) -> None:
         """Embed all images in a folder and save them in a .npy file
         """
+        assert output_path.endswith(".npy"), "`output_path` must end with .npy"
         embeddings = {}
-        for image in os.listdir(folder_path):
+        for image in tqdm(os.listdir(folder_path)):
             image_path = os.path.join(folder_path, image)
             image = Image.open(image_path)
             embedding = self.embed(image)
             embeddings[image] = embedding
-        np.save(folder_path + ".npy", embeddings)
+        np.save(output_path, embeddings)

geoguessr_bot/retriever/dino_v2_embedder.py

@@ -0,0 +1,31 @@
+import numpy as np
+import torch
+import torchvision.transforms as T
+from PIL import Image
+
+from .abstract_embedder import AbstractImageEmbedder
+
+
+class DinoV2Embedder(AbstractImageEmbedder):
+    def __init__(self, device: str = "cpu"):
+        """Embedder using DINOv2 embeddings.
+        """
+        super().__init__(device)
+        self.model = torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14').to(self.device)
+        self.model.eval()
+        self.transforms = T.Compose([
+            T.Resize((256, 256), interpolation=T.InterpolationMode.BICUBIC),
+            T.CenterCrop(224),
+            T.ToTensor(),
+            T.Normalize(
+                mean=[0.485, 0.456, 0.406],
+                std=[0.229, 0.224, 0.225]
+            )
+        ])
+
+    def embed(self, image: Image) -> np.ndarray:
+        image = image.convert("RGB")
+        image = self.transforms(image).unsqueeze(0).to(self.device)
+        with torch.no_grad():
+            output = self.model(image)[0].cpu().numpy()
+        return output

geoguessr_bot/retriever/retriever.py

@@ -20,11 +20,11 @@ class Retriever:
 
     @staticmethod
     def load_embeddings(embeddings_path: str) -> Dict[str, np.ndarray]:
-        """Load embeddings from a file
+        """Load embeddings from a .npy file
         """
-        raise NotImplementedError
+        return np.load(embeddings_path, allow_pickle=True).item()
 
-    def retrieve(self, queries: np.ndarray, n_neighbors: int = 5) -> List[str]:
+    def retrieve(self, queries: np.ndarray, n_neighbors: int = 5) -> List[List[str]]:
         """Retrieve nearest neighbors indexes from queries
         """
         dist, indexes = self.index.search(queries, n_neighbors)

requirements.txt

@@ -7,3 +7,8 @@ Pillow
 transformers
 requests
 faiss-cpu
+torch
+torchvision
+tqdm
+configue
+fire