geoguessr_bot/retriever/abstract_embedder.py

import os
from PIL import Image
import numpy as np
from tqdm import tqdm


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, output_path: str) -> None:
        """Embed all images in a folder and save them in a .npy file
        """
        assert output_path.endswith(".npy"), "`output_path` must end with .npy"
        embeddings = {}
        for name in tqdm(os.listdir(folder_path)):
            image_path = os.path.join(folder_path, name)
            image = Image.open(image_path)
            embedding = self.embed(image)
            embeddings[name] = embedding
        np.save(output_path, embeddings)