faiss.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303
  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 ..base import BaseComponent
  21. class FaissIndexer(BaseComponent):
  22. INPUT_KEYS = "feature"
  23. OUTPUT_KEYS = ["label", "score"]
  24. DEAULT_INPUTS = {"feature": "feature"}
  25. DEAULT_OUTPUTS = {"label": "label", "score": "score"}
  26. ENABLE_BATCH = True
  27. def __init__(
  28. self,
  29. index_bytes=None,
  30. vector_path=None,
  31. id_map=None,
  32. metric_type="IP",
  33. return_k=1,
  34. score_thres=None,
  35. hamming_radius=None,
  36. ):
  37. super().__init__()
  38. if metric_type == "hamming":
  39. if index_bytes is not None:
  40. self._indexer = faiss.deserialize_index(index_bytes)
  41. else:
  42. self._indexer = faiss.read_index_binary(vector_path)
  43. self.hamming_radius = hamming_radius
  44. else:
  45. if index_bytes is not None:
  46. self._indexer = faiss.deserialize_index(index_bytes)
  47. else:
  48. self._indexer = faiss.read_index(vector_path)
  49. self.score_thres = score_thres
  50. self.id_map = id_map
  51. self.metric_type = metric_type
  52. self.return_k = return_k
  53. def apply(self, feature):
  54. """apply"""
  55. scores_list, ids_list = self._indexer.search(np.array(feature), self.return_k)
  56. preds = []
  57. for scores, ids in zip(scores_list, ids_list):
  58. labels = []
  59. for id in ids:
  60. if id > 0:
  61. labels.append(self.id_map[id])
  62. preds.append({"score": scores, "label": labels})
  63. if self.metric_type == "hamming":
  64. idxs = np.where(scores_list[:, 0] > self.hamming_radius)[0]
  65. else:
  66. idxs = np.where(scores_list[:, 0] < self.score_thres)[0]
  67. for idx in idxs:
  68. preds[idx] = {"score": None, "label": None}
  69. return preds
  70. class FaissBuilder:
  71. SUPPORT_MODE = ("new", "remove", "append")
  72. SUPPORT_METRIC_TYPE = ("hamming", "IP", "L2")
  73. SUPPORT_INDEX_TYPE = ("Flat", "IVF", "HNSW32")
  74. BINARY_METRIC_TYPE = ("hamming",)
  75. BINARY_SUPPORT_INDEX_TYPE = ("Flat", "IVF", "BinaryHash")
  76. def __init__(self, predict, mode="new", index_type="HNSW32", metric_type="IP"):
  77. super().__init__()
  78. assert (
  79. mode in self.SUPPORT_MODE
  80. ), f"Supported modes only: {self.SUPPORT_MODE}. But received {mode}!"
  81. assert (
  82. metric_type in self.SUPPORT_METRIC_TYPE
  83. ), f"Supported metric types only: {self.SUPPORT_METRIC_TYPE}!"
  84. assert (
  85. index_type in self.SUPPORT_INDEX_TYPE
  86. ), f"Supported index types only: {self.SUPPORT_INDEX_TYPE}!"
  87. self._predict = predict
  88. self._mode = mode
  89. self._metric_type = metric_type
  90. self._index_type = index_type
  91. def _get_index_type(self, num=None):
  92. # if IVF method, cal ivf number automaticlly
  93. if self._index_type == "IVF":
  94. index_type = self._index_type + str(min(int(num // 8), 65536))
  95. if self._metric_type in self.BINARY_METRIC_TYPE:
  96. index_type += ",BFlat"
  97. else:
  98. index_type += ",Flat"
  99. # for binary index, add B at head of index_type
  100. if self._metric_type in self.BINARY_METRIC_TYPE:
  101. assert (
  102. self._index_type in self.BINARY_SUPPORT_INDEX_TYPE
  103. ), f"The metric type({self._metric_type}) only support {self.BINARY_SUPPORT_INDEX_TYPE} index types!"
  104. index_type = "B" + index_type
  105. if self._index_type == "HNSW32":
  106. logging.warning("The HNSW32 method dose not support 'remove' operation")
  107. index_type = "HNSW32"
  108. if self._index_type == "Flat":
  109. index_type = "Flat"
  110. return index_type
  111. def _get_metric_type(self):
  112. if self._metric_type == "hamming":
  113. return faiss.METRIC_Hamming
  114. elif self._metric_type == "jaccard":
  115. return faiss.METRIC_Jaccard
  116. elif self._metric_type == "IP":
  117. return faiss.METRIC_INNER_PRODUCT
  118. elif self._metric_type == "L2":
  119. return faiss.METRIC_L2
  120. def build(
  121. self,
  122. label_file,
  123. image_root,
  124. index_dir=None,
  125. ):
  126. file_list, gallery_docs = get_file_list(label_file, image_root)
  127. features = [res["feature"] for res in self._predict(file_list)]
  128. dtype = np.uint8 if self._metric_type in self.BINARY_METRIC_TYPE else np.float32
  129. features = np.array(features).astype(dtype)
  130. vector_num, vector_dim = features.shape
  131. if self._metric_type in self.BINARY_METRIC_TYPE:
  132. index = faiss.index_binary_factory(
  133. vector_dim,
  134. self._get_index_type(vector_num),
  135. self._get_metric_type(),
  136. )
  137. else:
  138. index = faiss.index_factory(
  139. vector_dim,
  140. self._get_index_type(vector_num),
  141. self._get_metric_type(),
  142. )
  143. index = faiss.IndexIDMap2(index)
  144. ids = {}
  145. # calculate id for new data
  146. index, ids = self._add_gallery(index, ids, features, gallery_docs)
  147. if index_dir:
  148. self._save_gallery(index, ids, index_dir)
  149. return faiss.serialize_index(index), ids
  150. def remove(
  151. self,
  152. label_file,
  153. image_root,
  154. index_dir=None,
  155. index_bytes=None,
  156. vector_path=None,
  157. id_map=None,
  158. ):
  159. file_list, gallery_docs = get_file_list(label_file, image_root)
  160. if index_bytes is not None:
  161. index = faiss.deserialize_index(index_bytes)
  162. ids = id_map
  163. else:
  164. # load vector.index and id_map.pkl
  165. index, ids = self._load_index(index_dir)
  166. if self._index_type == "HNSW32":
  167. raise RuntimeError(
  168. "The index_type: HNSW32 dose not support 'remove' operation"
  169. )
  170. # remove ids in id_map, remove index data in faiss index
  171. index, ids = self._rm_id_in_galllery(index, ids, gallery_docs)
  172. if index_dir:
  173. self._save_gallery(index, ids, index_dir)
  174. return faiss.serialize_index(index), ids
  175. def append(
  176. self,
  177. label_file,
  178. image_root,
  179. index_dir=None,
  180. index_bytes=None,
  181. vector_path=None,
  182. id_map=None,
  183. ):
  184. file_list, gallery_docs = get_file_list(label_file, image_root)
  185. features = [res["feature"] for res in self._predict(file_list)]
  186. dtype = np.uint8 if self._metric_type in self.BINARY_METRIC_TYPE else np.float32
  187. features = np.array(features).astype(dtype)
  188. if index_bytes is not None:
  189. index = faiss.deserialize_index(index_bytes)
  190. ids = id_map
  191. else:
  192. # load vector.index and id_map.pkl
  193. index, ids = self._load_index(index_dir)
  194. # calculate id for new data
  195. index, ids = self._add_gallery(index, ids, features, gallery_docs)
  196. if index_dir:
  197. self._save_gallery(index, ids, index_dir)
  198. return faiss.serialize_index(index), ids
  199. def _load_index(self, index_dir):
  200. assert os.path.join(
  201. index_dir, "vector.index"
  202. ), "The vector.index dose not exist in {} when 'index_operation' is not None".format(
  203. index_dir
  204. )
  205. assert os.path.join(
  206. index_dir, "id_map.pkl"
  207. ), "The id_map.pkl dose not exist in {} when 'index_operation' is not None".format(
  208. index_dir
  209. )
  210. index = faiss.read_index(os.path.join(index_dir, "vector.index"))
  211. with open(os.path.join(index_dir, "id_map.pkl"), "rb") as fd:
  212. ids = pickle.load(fd)
  213. assert index.ntotal == len(
  214. ids.keys()
  215. ), "data number in index is not equal in in id_map"
  216. return index, ids
  217. def _add_gallery(self, index, ids, gallery_features, gallery_docs):
  218. start_id = max(ids.keys()) + 1 if ids else 0
  219. ids_now = (np.arange(0, len(gallery_docs)) + start_id).astype(np.int64)
  220. # only train when new index file
  221. if self._mode == "new":
  222. if self._metric_type in self.BINARY_METRIC_TYPE:
  223. index.add(gallery_features)
  224. else:
  225. index.train(gallery_features)
  226. if not self._metric_type in self.BINARY_METRIC_TYPE:
  227. index.add_with_ids(gallery_features, ids_now)
  228. for i, d in zip(list(ids_now), gallery_docs):
  229. ids[i] = d
  230. return index, ids
  231. def _rm_id_in_galllery(self, index, ids, gallery_docs):
  232. remove_ids = list(filter(lambda k: ids.get(k) in gallery_docs, ids.keys()))
  233. remove_ids = np.asarray(remove_ids)
  234. index.remove_ids(remove_ids)
  235. for k in remove_ids:
  236. del ids[k]
  237. return index, ids
  238. def _save_gallery(self, index, ids, index_dir):
  239. Path(index_dir).mkdir(parents=True, exist_ok=True)
  240. if self._metric_type in self.BINARY_METRIC_TYPE:
  241. faiss.write_index_binary(index, os.path.join(index_dir, "vector.index"))
  242. else:
  243. faiss.write_index(index, os.path.join(index_dir, "vector.index"))
  244. with open(os.path.join(index_dir, "id_map.pkl"), "wb") as fd:
  245. pickle.dump(ids, fd)
  246. def get_file_list(data_file, root_dir, delimiter=" "):
  247. root_dir = Path(root_dir)
  248. files = []
  249. labels = []
  250. lines = []
  251. with open(data_file, "r", encoding="utf-8") as f:
  252. lines = f.readlines()
  253. for line in lines:
  254. path, label = line.strip().split(delimiter)
  255. file_path = root_dir / path
  256. files.append(file_path.as_posix())
  257. labels.append(label)
  258. return files, labels