src/faiss_db.py

from functools import lru_cache
import faiss
import json
import numpy as np


class FAISS_DB:
    """
    A class for interacting with a FAISS index and keyframes group information.

    Args:
        faiss_index_file_path (str): The path to the FAISS index file.
        keyframes_groups_json_path (str): The path to the JSON file containing keyframes group information.

    Attributes:
        faiss_index_file_path (str): The path to the FAISS index file.
        keyframes_groups_json_path (str): The path to the JSON file containing keyframes group information.
        index (faiss.Index): The FAISS index, loaded lazily.
        keyframes_group_info (dict): The keyframes group information, loaded lazily.

    Example:
        >>> db = FAISS_DB('faiss_index.bin', 'keyframes_groups.json')
    """

    def __init__(
        self, faiss_index_file_path: str, keyframes_groups_json_path: str
    ) -> None:
        """
        Initialize the FAISS_DB instance.

        Args:
            faiss_index_file_path (str): The path to the FAISS index file.
            keyframes_groups_json_path (str): The path to the JSON file containing keyframes group information.
        """
        self.faiss_index_file_path = faiss_index_file_path
        self.keyframes_groups_json_path = keyframes_groups_json_path

    @property
    @lru_cache(maxsize=1)
    def index(self) -> faiss.Index:
        """
        The FAISS index, loaded lazily.

        Returns:
            faiss.Index: The FAISS index.

        Example:
            >>> db.index
        """
        return faiss.read_index(self.faiss_index_file_path)

    @property
    @lru_cache(maxsize=1)
    def keyframes_group_info(self) -> dict[str, str]:
        """
        The keyframes group information, loaded lazily.

        Returns:
            dict: The keyframes group information.

        Example:
            >>> db.keyframes_group_info
        """
        with open(self.keyframes_groups_json_path) as file:
            return json.load(file)

    def kNN_search(
        self, query_vector: np.ndarray, topk: int = 10
    ) -> list[dict[str, str]]:
        """
        Perform a k-NN search on the FAISS index.

        Args:
            query_vector (np.ndarray): The query vector.
            topk (int, optional): The number of nearest neighbors to return. Defaults to 10.

        Returns:
            list[dict[str, str]]: A list of dictionaries, 
                where each dictionary contains the distance and keyframes group information for a nearest neighbor.

        Example:
            >>> db = FAISS_DB('faiss_index.bin', 'keyframes_groups.json')
            >>> query_vector = np.random.rand(1, 768)  # Replace with your query vector
            >>> results = db.kNN_search(query_vector)
            >>> print(results)  # Print the k-NN search results
        """
        distances, ids = self.index.search(query_vector, topk)
        distances, ids = distances[0], ids[0]
        return [
            {"distance": str(dist), **self.keyframes_group_info[int(ids_i)]}
            for dist, ids_i in zip(distances, ids)
        ]