Returns distances and indices of nearest neighbour search. Args: query_features: Features to retrieve. index_features: [optional] Index features to search in.
(
self,
n_nearest_neighbours,
query_features: np.ndarray,
index_features: np.ndarray = None,
)
| 64 | pass |
| 65 | |
| 66 | def run( |
| 67 | self, |
| 68 | n_nearest_neighbours, |
| 69 | query_features: np.ndarray, |
| 70 | index_features: np.ndarray = None, |
| 71 | ) -> Union[np.ndarray, np.ndarray, np.ndarray]: |
| 72 | """ |
| 73 | Returns distances and indices of nearest neighbour search. |
| 74 | |
| 75 | Args: |
| 76 | query_features: Features to retrieve. |
| 77 | index_features: [optional] Index features to search in. |
| 78 | """ |
| 79 | if index_features is None: |
| 80 | return self.search_index.search(query_features, n_nearest_neighbours) |
| 81 | |
| 82 | # Build a search index just for this search. |
| 83 | search_index = self._create_index(index_features.shape[-1]) |
| 84 | self._train(search_index, index_features) |
| 85 | search_index.add(index_features) |
| 86 | return search_index.search(query_features, n_nearest_neighbours) |
| 87 | |
| 88 | def save(self, filename: str) -> None: |
| 89 | faiss.write_index(self._index_to_cpu(self.search_index), filename) |