salut

geoguessr_bot/guessr/global_embedder_guessr.py

@@ -1,4 +1,6 @@
+from dataclasses import dataclass
 from geoguessr_bot.guessr import AbstractGuessr
+from geoguessr_bot.interfaces import Coordinate
 from geoguessr_bot.retriever import AbstractImageEmbedder
 from geoguessr_bot.retriever import Retriever
 

geoguessr_bot/retriever/abstract_embedder.py

@@ -1,3 +1,8 @@
+import os
+from PIL import Image
+import numpy as np
+
+
 class AbstractImageEmbedder:
     def __init__(self, device: str = "cpu"):
         self.device = device

geoguessr_bot/retriever/dino_embedder.py

@@ -1,3 +1,5 @@
+import numpy as np
+from PIL import Image
 from transformers import ViTFeatureExtractor, ViTModel
 
 from .abstract_embedder import AbstractImageEmbedder
@@ -10,9 +12,9 @@ class DinoEmbedder(AbstractImageEmbedder):
         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")
+        inputs = self.feature_extractor(images=image, return_tensors="pt")
         for key in inputs:
             inputs[key] = inputs[key].to(self.device)
-        outputs = model(**inputs)
+        outputs = self.model(**inputs)
         last_hidden_states = outputs.last_hidden_state.to("cpu").numpy()
         return last_hidden_states

geoguessr_bot/retriever/retriever.py

@@ -1,3 +1,9 @@
+from typing import Dict, List
+
+import numpy as np
+import faiss
+
+
 class Retriever:
     def __init__(self, embeddings_path: str):
         self.embeddings: Dict[str, np.ndarray] = self.load_embeddings(embeddings_path)