geoguessr_bot/guessr/global_embedder_guessr.py

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


@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