faiss.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349
  1. # copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. import os
  15. import pickle
  16. from pathlib import Path
  17. import faiss
  18. import numpy as np
  19. from ....utils import logging
  20. from ...utils.io import YAMLWriter, YAMLReader
  21. from ..base import BaseComponent
  22. class IndexData:
  23. VECTOR_FN = "vector"
  24. VECTOR_SUFFIX = ".index"
  25. IDMAP_FN = "id_map"
  26. IDMAP_SUFFIX = ".yaml"
  27. def __init__(self, index, id_map, metric_type, index_type):
  28. self._index = index
  29. self._id_map = id_map
  30. self._metric_type = metric_type
  31. self._index_type = index_type
  32. @property
  33. def index(self):
  34. return self._index
  35. @property
  36. def index_bytes(self):
  37. return faiss.serialize_index(self._index)
  38. @property
  39. def id_map(self):
  40. return self._id_map
  41. @property
  42. def metric_type(self):
  43. return self._metric_type
  44. @property
  45. def index_type(self):
  46. return self._index_type
  47. @property
  48. def index_info(self):
  49. return {
  50. "index_type": self.index_type,
  51. "metric_type": self.metric_type,
  52. "id_map": self._convert_int(self.id_map),
  53. }
  54. def _convert_int(self, id_map):
  55. return {int(k): str(v) for k, v in id_map.items()}
  56. @staticmethod
  57. def _convert_int64(id_map):
  58. return {np.int64(k): str(v) for k, v in id_map.items()}
  59. def save(self, save_dir):
  60. save_dir = Path(save_dir)
  61. save_dir.mkdir(parents=True, exist_ok=True)
  62. vector_path = (save_dir / f"{self.VECTOR_FN}{self.VECTOR_SUFFIX}").as_posix()
  63. index_info_path = (save_dir / f"{self.IDMAP_FN}{self.IDMAP_SUFFIX}").as_posix()
  64. if self.metric_type in FaissBuilder.BINARY_METRIC_TYPE:
  65. faiss.write_index_binary(self.index, vector_path)
  66. else:
  67. faiss.write_index(self.index, vector_path)
  68. yaml_writer = YAMLWriter()
  69. yaml_writer.write(
  70. index_info_path,
  71. self.index_info,
  72. default_flow_style=False,
  73. allow_unicode=True,
  74. )
  75. @classmethod
  76. def load(cls, index):
  77. if isinstance(index, str):
  78. index_root = Path(index)
  79. vector_path = index_root / f"{cls.VECTOR_FN}{cls.VECTOR_SUFFIX}"
  80. index_info_path = index_root / f"{cls.IDMAP_FN}{cls.IDMAP_SUFFIX}"
  81. assert (
  82. vector_path.exists()
  83. ), f"Not found the {cls.VECTOR_FN}{cls.VECTOR_SUFFIX} file in {index}!"
  84. assert (
  85. index_info_path.exists()
  86. ), f"Not found the {cls.IDMAP_FN}{cls.IDMAP_SUFFIX} file in {index}!"
  87. yaml_reader = YAMLReader()
  88. index_info = yaml_reader.read(index_info_path)
  89. assert (
  90. "id_map" in index_info
  91. and "metric_type" in index_info
  92. and "index_type" in index_info
  93. ), f"The index_info file({index_info_path}) may have been damaged, `id_map` or `metric_type` or `index_type` not found in `index_info`."
  94. id_map = IndexData._convert_int64(index_info["id_map"])
  95. if index_info["metric_type"] in FaissBuilder.BINARY_METRIC_TYPE:
  96. index = faiss.read_index_binary(vector_path.as_posix())
  97. else:
  98. index = faiss.read_index(vector_path.as_posix())
  99. assert index.ntotal == len(
  100. id_map
  101. ), "data number in index is not equal in in id_map"
  102. return index, id_map, index_info["metric_type"], index_info["index_type"]
  103. else:
  104. assert isinstance(index, IndexData)
  105. return index.index, index.id_map, index.metric_type, index.index_type
  106. class FaissIndexer(BaseComponent):
  107. INPUT_KEYS = "feature"
  108. OUTPUT_KEYS = ["label", "score"]
  109. DEAULT_INPUTS = {"feature": "feature"}
  110. DEAULT_OUTPUTS = {"label": "label", "score": "score"}
  111. ENABLE_BATCH = True
  112. def __init__(
  113. self,
  114. index,
  115. return_k=1,
  116. score_thres=None,
  117. hamming_radius=None,
  118. ):
  119. super().__init__()
  120. self._indexer, self.id_map, self.metric_type, index_type = IndexData.load(index)
  121. self.return_k = return_k
  122. if self.metric_type in FaissBuilder.BINARY_METRIC_TYPE:
  123. self.hamming_radius = hamming_radius
  124. else:
  125. self.score_thres = score_thres
  126. def apply(self, feature):
  127. """apply"""
  128. scores_list, ids_list = self._indexer.search(np.array(feature), self.return_k)
  129. preds = []
  130. for scores, ids in zip(scores_list, ids_list):
  131. labels = []
  132. for id in ids:
  133. if id > 0:
  134. labels.append(self.id_map[id])
  135. preds.append({"score": scores, "label": labels})
  136. if self.metric_type in FaissBuilder.BINARY_METRIC_TYPE:
  137. idxs = np.where(scores_list[:, 0] > self.hamming_radius)[0]
  138. else:
  139. idxs = np.where(scores_list[:, 0] < self.score_thres)[0]
  140. for idx in idxs:
  141. preds[idx] = {"score": None, "label": None}
  142. return preds
  143. class FaissBuilder:
  144. SUPPORT_METRIC_TYPE = ("hamming", "IP", "L2")
  145. SUPPORT_INDEX_TYPE = ("Flat", "IVF", "HNSW32")
  146. BINARY_METRIC_TYPE = ("hamming",)
  147. BINARY_SUPPORT_INDEX_TYPE = ("Flat", "IVF", "BinaryHash")
  148. @classmethod
  149. def _get_index_type(cls, metric_type, index_type, num=None):
  150. # if IVF method, cal ivf number automaticlly
  151. if index_type == "IVF":
  152. index_type = index_type + str(min(int(num // 8), 65536))
  153. if metric_type in cls.BINARY_METRIC_TYPE:
  154. index_type += ",BFlat"
  155. else:
  156. index_type += ",Flat"
  157. # for binary index, add B at head of index_type
  158. if metric_type in cls.BINARY_METRIC_TYPE:
  159. assert (
  160. index_type in cls.BINARY_SUPPORT_INDEX_TYPE
  161. ), f"The metric type({metric_type}) only support {cls.BINARY_SUPPORT_INDEX_TYPE} index types!"
  162. index_type = "B" + index_type
  163. if index_type == "HNSW32":
  164. logging.warning("The HNSW32 method dose not support 'remove' operation")
  165. index_type = "HNSW32"
  166. if index_type == "Flat":
  167. index_type = "Flat"
  168. return index_type
  169. @classmethod
  170. def _get_metric_type(cls, metric_type):
  171. if metric_type == "hamming":
  172. return faiss.METRIC_Hamming
  173. elif metric_type == "jaccard":
  174. return faiss.METRIC_Jaccard
  175. elif metric_type == "IP":
  176. return faiss.METRIC_INNER_PRODUCT
  177. elif metric_type == "L2":
  178. return faiss.METRIC_L2
  179. @classmethod
  180. def build(
  181. cls,
  182. gallery_imgs,
  183. gallery_label,
  184. predict_func,
  185. metric_type="IP",
  186. index_type="HNSW32",
  187. ):
  188. assert (
  189. index_type in cls.SUPPORT_INDEX_TYPE
  190. ), f"Supported index types only: {cls.SUPPORT_INDEX_TYPE}!"
  191. assert (
  192. metric_type in cls.SUPPORT_METRIC_TYPE
  193. ), f"Supported metric types only: {cls.SUPPORT_METRIC_TYPE}!"
  194. if isinstance(gallery_label, str):
  195. gallery_docs, gallery_list = cls.load_gallery(gallery_label, gallery_imgs)
  196. else:
  197. gallery_docs, gallery_list = gallery_label, gallery_imgs
  198. features = [res["feature"] for res in predict_func(gallery_list)]
  199. dtype = np.uint8 if metric_type in cls.BINARY_METRIC_TYPE else np.float32
  200. features = np.array(features).astype(dtype)
  201. vector_num, vector_dim = features.shape
  202. if metric_type in cls.BINARY_METRIC_TYPE:
  203. index = faiss.index_binary_factory(
  204. vector_dim,
  205. cls._get_index_type(metric_type, index_type, vector_num),
  206. cls._get_metric_type(metric_type),
  207. )
  208. else:
  209. index = faiss.index_factory(
  210. vector_dim,
  211. cls._get_index_type(metric_type, index_type, vector_num),
  212. cls._get_metric_type(metric_type),
  213. )
  214. index = faiss.IndexIDMap2(index)
  215. ids = {}
  216. # calculate id for new data
  217. index, ids = cls._add_gallery(
  218. metric_type, index, ids, features, gallery_docs, mode="new"
  219. )
  220. return IndexData(index, ids, metric_type, index_type)
  221. @classmethod
  222. def remove(
  223. cls,
  224. remove_ids,
  225. index,
  226. ):
  227. index, ids, metric_type, index_type = IndexData.load(index)
  228. if index_type == "HNSW32":
  229. raise RuntimeError(
  230. "The index_type: HNSW32 dose not support 'remove' operation"
  231. )
  232. if isinstance(remove_ids, str):
  233. lines = []
  234. with open(remove_ids) as f:
  235. lines = f.readlines()
  236. remove_ids = []
  237. for line in lines:
  238. id_ = int(line.strip().split(" ")[0])
  239. remove_ids.append(id_)
  240. remove_ids = np.asarray(remove_ids)
  241. else:
  242. remove_ids = np.asarray(remove_ids)
  243. # remove ids in id_map, remove index data in faiss index
  244. index.remove_ids(remove_ids)
  245. ids = {k: v for k, v in ids.items() if k not in remove_ids}
  246. return IndexData(index, ids, metric_type, index_type)
  247. @classmethod
  248. def append(cls, gallery_imgs, gallery_label, predict_func, index):
  249. index, ids, metric_type, index_type = IndexData.load(index)
  250. assert (
  251. metric_type in cls.SUPPORT_METRIC_TYPE
  252. ), f"Supported metric types only: {cls.SUPPORT_METRIC_TYPE}!"
  253. if isinstance(gallery_label, str):
  254. gallery_docs, gallery_list = cls.load_gallery(gallery_label, gallery_imgs)
  255. else:
  256. gallery_docs, gallery_list = gallery_label, gallery_imgs
  257. features = [res["feature"] for res in predict_func(gallery_list)]
  258. dtype = np.uint8 if metric_type in cls.BINARY_METRIC_TYPE else np.float32
  259. features = np.array(features).astype(dtype)
  260. # calculate id for new data
  261. index, ids = cls._add_gallery(
  262. metric_type, index, ids, features, gallery_docs, mode="append"
  263. )
  264. return IndexData(index, ids, metric_type, index_type)
  265. @classmethod
  266. def _add_gallery(
  267. cls, metric_type, index, ids, gallery_features, gallery_docs, mode
  268. ):
  269. start_id = max(ids.keys()) + 1 if ids else 0
  270. ids_now = (np.arange(0, len(gallery_docs)) + start_id).astype(np.int64)
  271. # only train when new index file
  272. if mode == "new":
  273. if metric_type in cls.BINARY_METRIC_TYPE:
  274. index.add(gallery_features)
  275. else:
  276. index.train(gallery_features)
  277. if not metric_type in cls.BINARY_METRIC_TYPE:
  278. index.add_with_ids(gallery_features, ids_now)
  279. for i, d in zip(list(ids_now), gallery_docs):
  280. ids[i] = d
  281. return index, ids
  282. @classmethod
  283. def load_gallery(cls, gallery_label_path, gallery_imgs_root="", delimiter=" "):
  284. lines = []
  285. files = []
  286. labels = []
  287. root = Path(gallery_imgs_root)
  288. with open(gallery_label_path, "r", encoding="utf-8") as f:
  289. lines = f.readlines()
  290. for line in lines:
  291. path, label = line.strip().split(delimiter)
  292. file_path = root / path
  293. files.append(file_path.as_posix())
  294. labels.append(label)
  295. return labels, files