geoguessr_bot/retriever/abstract_embedder.py

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)