Spaces:
Sleeping
Sleeping
File size: 3,112 Bytes
84ee555 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | 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)
]
|