wip embedder

app.py

@@ -5,7 +5,7 @@ import plotly.graph_objects as go
 from geoguessr_bot.guessr import RandomGuessr, AbstractGuessr
 
 ALL_GUESSR_CLASS = {
-    "random": RandomGuessr
+    "random": RandomGuessr,
 }
 
 ALL_GUESSR_ARGS = {

geoguessr_bot/guessr/__init__.py

@@ -1,2 +1,3 @@
 from .abstract_guessr import AbstractGuessr
 from .random_guessr import RandomGuessr
+from .global_embedder_guessr import GlobalEmbedderGuessr

geoguessr_bot/guessr/global_embedder_guessr.py

@@ -0,0 +1,25 @@
+from geoguessr_bot.guessr import AbstractGuessr
+from geoguessr_bot.retriever import AbstractImageEmbedder
+from geoguessr_bot.retriever import Retriever
+
+
+@dataclass
+class GlobalEmbedderGuessr(AbstractGuessr):
+    """Guesses a coordinate using an Embedder and a retriever
+    """
+
+    embedder: AbstractImageEmbedder
+    retriever: Retriever
+
+    def guess(self, image: Image) -> Coordinate:
+        """Guess a coordinate from an image
+        """
+        # Embed image
+        image_embedding = self.embedder.embed(image)
+
+        # Retrieve nearest neighbors
+        nearest_neighbors = self.retriever.retrieve(image_embedding)
+
+        # Guess coordinate
+        guess_coordinate = self.retriever.image_to_coordinate[nearest_neighbors[0][0]]
+        return guess_coordinate

geoguessr_bot/retriever/__init__.py

@@ -0,0 +1,3 @@
+from .retriever import AbstractRetriever
+from .abstract_embedder import AbstractImageEmbedder
+from .dino_embedder import DinoEmbedder

geoguessr_bot/retriever/abstract_embedder.py

@@ -0,0 +1,19 @@
+class AbstractImageEmbedder:
+    def __init__(self, device: str = "cpu"):
+        self.device = device
+
+    def embed(self, image: Image) -> np.ndarray:
+        """Embed an image
+        """
+        raise NotImplementedError
+
+    def embed_folder(self, folder_path: str):
+        """Embed all images in a folder and save them in a .npy file
+        """
+        embeddings = {}
+        for image in 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)

geoguessr_bot/retriever/dino_embedder.py

@@ -0,0 +1,18 @@
+from transformers import ViTFeatureExtractor, ViTModel
+
+from .abstract_embedder import AbstractImageEmbedder
+
+
+class DinoEmbedder(AbstractImageEmbedder):
+    def __init__(self, device: str = "cpu", model_name: str = "facebook/dino-vitb8"):
+        super().__init__(device)
+        self.feature_extractor = ViTFeatureExtractor.from_pretrained(model_name)
+        self.model = ViTModel.from_pretrained(model_name).to(self.device)
+
+    def embed(self, image: Image) -> np.ndarray:
+        inputs = feature_extractor(images=image, return_tensors="pt")
+        for key in inputs:
+            inputs[key] = inputs[key].to(self.device)
+        outputs = model(**inputs)
+        last_hidden_states = outputs.last_hidden_state.to("cpu").numpy()
+        return last_hidden_states

geoguessr_bot/retriever/retriever.py

@@ -0,0 +1,26 @@
+class Retriever:
+    def __init__(self, embeddings_path: str):
+        self.embeddings: Dict[str, np.ndarray] = self.load_embeddings(embeddings_path)
+
+        # Keep track of image names
+        self.image_to_index = {image_name: i for i, image_name in enumerate(self.embeddings.keys())}
+        self.index_to_image = {i: image_name for i, image_name in enumerate(self.embeddings.keys())}
+
+        # Build Faiss index
+        self.embeddings = np.array(list(self.embeddings.values()))
+        self.dim = self.embeddings.shape[1]
+        self.index = faiss.IndexFlatL2(self.dim)
+        self.index.add(self.embeddings)
+
+    @staticmethod
+    def load_embeddings(embeddings_path: str) -> Dict[str, np.ndarray]:
+        """Load embeddings from a file
+        """
+        raise NotImplementedError
+
+    def retrieve(self, queries: np.ndarray, n_neighbors: int = 5) -> List[str]:
+        """Retrieve nearest neighbors indexes from queries
+        """
+        dist, indexes = self.index.search(queries, n_neighbors)
+        return [[self.index_to_image[i] for i in index] for index in indexes]
+

requirements.txt

@@ -4,3 +4,6 @@ plotly
 gradio
 pydantic
 Pillow
+transformers
+requests
+faiss-cpu