atiso-clip / src /itr /dtb_cursor.py
ngxquang
feat: update clip api for deployment
2e5d83b
raw
history blame contribute delete
991 Bytes
import json
from functools import lru_cache
import faiss
class DatabaseCursor:
def __init__(self, index_file_path: str, keyframes_groups_json_path: str):
self._load_index(index_file_path)
self._load_keyframes_groups_info(keyframes_groups_json_path)
@lru_cache(maxsize=1)
def _load_index(self, index_file_path):
self.index = faiss.read_index(index_file_path)
@lru_cache(maxsize=1)
def _load_keyframes_groups_info(self, keyframes_groups_json_path: str):
with open(keyframes_groups_json_path) as file:
self.keyframes_group_info = json.loads(file.read())
def kNN_search(self, query_vector: str, topk: int = 10):
results = []
distances, ids = self.index.search(query_vector, topk)
for i in range(len(ids[0])):
frame_detail = self.keyframes_group_info[ids[0][i]]
frame_detail["distance"] = str(distances[0][i])
results.append(frame_detail)
return results