[ADD] Average embedder

app.py

@@ -4,18 +4,20 @@ import numpy as np
 import gradio as gr
 import plotly.graph_objects as go
 
-from geoguessr_bot.guessr import RandomGuessr, AbstractGuessr, GlobalEmbedderGuessr
+from geoguessr_bot.guessr import RandomGuessr, AbstractGuessr, NearestNeighborEmbedderGuessr, \
+    AverageNeighborsEmbedderGuessr
 from geoguessr_bot.retriever import DinoV2Embedder, Retriever
 
 
 ALL_GUESSR_CLASS = {
     "random": RandomGuessr,
-    "globalEmbedder": GlobalEmbedderGuessr
+    "nearestNeighborEmbedder": NearestNeighborEmbedderGuessr,
+    "averageNeighborsEmbedder": AverageNeighborsEmbedderGuessr,
 }
 
 ALL_GUESSR_ARGS = {
     "random": {},
-    "globalEmbedder": {
+    "nearestNeighborEmbedder": {
         "embedder": DinoV2Embedder(
             device="cpu"
         ),
@@ -25,6 +27,19 @@ ALL_GUESSR_ARGS = {
         ),
         "metadata_path": os.path.join(os.path.dirname(os.path.abspath(__file__)),
                                       "resources/metadatav3.csv"),
+    },
+    "averageNeighborsEmbedder": {
+        "embedder": DinoV2Embedder(
+            device="cpu"
+        ),
+        "retriever": Retriever(
+            embeddings_path=os.path.join(os.path.dirname(os.path.abspath(__file__)),
+                                         "resources/embeddings.npy"),
+        ),
+        "metadata_path": os.path.join(os.path.dirname(os.path.abspath(__file__)),
+                                        "resources/metadatav3.csv"),
+        "n_neighbors": 2000,
+        "dbscan_eps": 0.5
     }
 }
 

geoguessr_bot/guessr/__init__.py

@@ -1,3 +1,4 @@
 from .abstract_guessr import AbstractGuessr
 from .random_guessr import RandomGuessr
-from .global_embedder_guessr import GlobalEmbedderGuessr
+from .nearest_neighbor_embedder_guessr import NearestNeighborEmbedderGuessr
+from .average_neighbor_embedder_guessr import AverageNeighborsEmbedderGuessr

geoguessr_bot/guessr/abstract_guessr.py

@@ -25,7 +25,7 @@ class AbstractGuessr:
         """Create an interactive map showing a coordinate
         """
         fig = go.Figure(go.Scattermapbox(
-            customdata=[guess_coordinate.__str__()] if guess_coordinate is not None else None,
+            customdata=[str(guess_coordinate)] if guess_coordinate is not None else None,
             lat=[guess_coordinate.latitude] if guess_coordinate is not None else None,
             lon=[guess_coordinate.longitude] if guess_coordinate is not None else None,
             mode="markers",

geoguessr_bot/guessr/average_neighbor_embedder_guessr.py

@@ -0,0 +1,64 @@
+from collections import Counter
+from dataclasses import dataclass
+
+import numpy as np
+from sklearn.cluster import DBSCAN
+from sklearn.metrics.pairwise import haversine_distances
+from PIL import Image
+import pandas as pd
+
+from geoguessr_bot.guessr import AbstractGuessr
+from geoguessr_bot.interfaces import Coordinate
+from geoguessr_bot.retriever import AbstractImageEmbedder
+from geoguessr_bot.retriever import Retriever
+
+
+@dataclass
+class AverageNeighborsEmbedderGuessr(AbstractGuessr):
+    """Guesses a coordinate using an Embedder and a retriever followed by NN.
+    """
+    embedder: AbstractImageEmbedder
+    retriever: Retriever
+    metadata_path: str
+    n_neighbors: int = 1000
+    dbscan_eps: float = 0.05
+
+    def __post_init__(self):
+        """Load metadata
+        """
+        metadata = pd.read_csv(self.metadata_path)
+        self.image_to_coordinate = {
+            image.split("/")[-1]: Coordinate(latitude=latitude, longitude=longitude)
+            for image, latitude, longitude in zip(metadata["path"], metadata["latitude"], metadata["longitude"])
+        }
+        # DBSCAN will be used to take the centroid of the biggest cluster among the N neighbors, using Haversine
+        self.dbscan = DBSCAN(eps=self.dbscan_eps, metric=haversine_distances)
+
+    def guess(self, image: Image) -> Coordinate:
+        """Guess a coordinate from an image
+        """
+        # Embed image
+        image = Image.fromarray(image)
+        image_embedding = self.embedder.embed(image)[None, :]
+        
+        # Retrieve nearest neighbors
+        nearest_neighbors, distances = self.retriever.retrieve(image_embedding, self.n_neighbors)
+        nearest_neighbors = nearest_neighbors[0]
+        distances = distances[0]
+
+        # Get coordinates of neighbors
+        neighbors_coordinates = [self.image_to_coordinate[nn].to_radians() for nn in nearest_neighbors]
+        neighbors_coordinates = np.array([[nn.latitude, nn.longitude] for nn in neighbors_coordinates])
+
+        # Use DBSCAN to find the biggest cluster and potentially remove outliers
+        clustering = self.dbscan.fit(neighbors_coordinates)
+        labels = clustering.labels_
+        biggest_cluster = max(Counter(labels))
+        neighbors_coordinates = neighbors_coordinates[labels == biggest_cluster]
+        distances = distances[labels == biggest_cluster]
+
+        # Guess coordinate as the closest image among the cluster regarding retrieving distance
+        guess_coordinate = neighbors_coordinates[np.argmin(distances)]
+        guess_coordinate = Coordinate.from_radians(guess_coordinate[0], guess_coordinate[1])
+        return guess_coordinate
+

geoguessr_bot/guessr/{global_embedder_guessr.py → nearest_neighbor_embedder_guessr.py}

@@ -10,10 +10,9 @@ from geoguessr_bot.retriever import Retriever
 
 
 @dataclass
-class GlobalEmbedderGuessr(AbstractGuessr):
-    """Guesses a coordinate using an Embedder and a retriever
+class NearestNeighborEmbedderGuessr(AbstractGuessr):
+    """Guesses a coordinate using an Embedder and a retriever followed by NN.
     """
-
     embedder: AbstractImageEmbedder
     retriever: Retriever
     metadata_path: str
@@ -35,9 +34,9 @@ class GlobalEmbedderGuessr(AbstractGuessr):
         image = Image.fromarray(image)
         image_embedding = self.embedder.embed(image)[None, :]
         
-        # Retrieve nearest neighbors
+        # Retrieve nearest neighbor
         nearest_neighbors = self.retriever.retrieve(image_embedding)
-        nearest_neighbor = nearest_neighbors[0][0]
+        nearest_neighbor = nearest_neighbors[0][0][0]
 
         # Guess coordinate
         guess_coordinate = self.image_to_coordinate[nearest_neighbor]

geoguessr_bot/interfaces.py

@@ -1,3 +1,4 @@
+import numpy as np
 from pydantic.main import BaseModel
 
 
@@ -7,3 +8,16 @@ class Coordinate(BaseModel):
 
     def __str__(self):
         return f"({round(self.latitude, 6)}, {round(self.longitude, 6)})"
+
+    def to_radians(self) -> 'Coordinate':
+        return Coordinate(
+            latitude=self.latitude * np.pi / 180.,
+            longitude=self.longitude * np.pi / 180.
+        )
+
+    @staticmethod
+    def from_radians(latitude: float, longitude: float) -> 'Coordinate':
+        return Coordinate(
+            latitude=latitude * 180. / np.pi,
+            longitude=longitude * 180. / np.pi
+        )

geoguessr_bot/retriever/retriever.py

@@ -1,13 +1,12 @@
-from typing import Dict, List
+from typing import Dict, List, Tuple
 
 import numpy as np
 import faiss
 
 
 class Retriever:
-    def __init__(self, embeddings_path: str, n_neighbors: int = 5):
+    def __init__(self, embeddings_path: str):
         self.embeddings: Dict[str, np.ndarray] = self.load_embeddings(embeddings_path)
-        self.n_neighbors = n_neighbors
 
         # Keep track of image names
         self.image_to_index = {image_name: i for i, image_name in enumerate(self.embeddings.keys())}
@@ -25,8 +24,8 @@ class Retriever:
         """
         return np.load(embeddings_path, allow_pickle=True).item()
 
-    def retrieve(self, queries: np.ndarray) -> List[List[str]]:
+    def retrieve(self, queries: np.ndarray, n_neighbors: int = 5) -> Tuple[List[List[str]], List[List[float]]]:
         """Retrieve nearest neighbors indexes from queries
         """
-        _, indexes = self.index.search(queries, self.n_neighbors)
-        return [[self.index_to_image[i] for i in index] for index in indexes]
+        distances, indexes = self.index.search(queries, n_neighbors)
+        return [[self.index_to_image[i] for i in index] for index in indexes], distances

requirements.txt

@@ -12,3 +12,4 @@ torchvision
 tqdm
 configue
 fire
+scikit-learn