refactor: use faiss instead of autofaiss (#80)

tti_eval/evaluation/image_retrieval.py

@@ -2,10 +2,10 @@
 from typing import Any
 
 import numpy as np
-from autofaiss import build_index
+from faiss import IndexFlatL2
 
 from tti_eval.common import Embeddings
-from tti_eval.utils import disable_tqdm, enable_tqdm
+from tti_eval.utils import disable_tqdm
 
 from .base import EvaluationModel
 
@@ -46,28 +46,29 @@ def __init__(
  self._class_counts[class_ids] = counts
 
  disable_tqdm() # Disable tqdm progress bar when building the index
- index, self.index_infos = build_index(self._val_embeddings.images, save_on_disk=False, verbose=logging.ERROR)
- enable_tqdm()
- if index is None:
- raise ValueError("Failed to build an index for knn search")
- self._index = index
-
- logger.info("knn classifier index_infos", extra=self.index_infos)
+ d = self._val_embeddings.images.shape[-1]
+ self._index = IndexFlatL2(d)
+ self._index.add(self._val_embeddings.images)
 
  def evaluate(self) -> float:
  _, nearest_indices = self._index.search(self._train_embeddings.images, self.k) # type: ignore
  nearest_classes = self._val_embeddings.labels[nearest_indices]
 
  # To compute retrieval accuracy, we ensure that a maximum of Q elements per sample are retrieved,
  # where Q represents the size of the respective class in the validation embeddings
- top_nearest_per_class = np.where(self._class_counts < self.k, self._class_counts, self.k)
+ top_nearest_per_class = np.where(
+ self._class_counts < self.k, self._class_counts, self.k
+ )
  top_nearest_per_sample = top_nearest_per_class[self._train_embeddings.labels]
 
  # Add a placeholder value for indices outside the retrieval scope
  nearest_classes[np.arange(self.k) >= top_nearest_per_sample[:, np.newaxis]] = -1
 
  # Count the number of neighbours that match the class of the sample and compute the mean accuracy
- matches_per_sample = np.sum(nearest_classes == np.array(self._train_embeddings.labels)[:, np.newaxis], axis=1)
+ matches_per_sample = np.sum(
+ nearest_classes == np.array(self._train_embeddings.labels)[:, np.newaxis],
+ axis=1,
+ )
  accuracies = np.divide(
  matches_per_sample,
  top_nearest_per_sample,

tti_eval/evaluation/knn.py

@@ -2,10 +2,10 @@
 from typing import Any
 
 import numpy as np
-from autofaiss import build_index
+from faiss import IndexFlatL2
 
 from tti_eval.common import ClassArray, Embeddings, ProbabilityArray
-from tti_eval.utils import disable_tqdm, enable_tqdm
+from tti_eval.utils import disable_tqdm
 
 from .base import ClassificationModel
 from .utils import softmax
@@ -44,15 +44,9 @@ def __init__(
  super().__init__(train_embeddings, validation_embeddings, num_classes)
  self.k = k
  disable_tqdm() # Disable tqdm progress bar when building the index
- index, self.index_infos = build_index(
- train_embeddings.images, metric_type="l2", save_on_disk=False, verbose=logging.ERROR
- )
- enable_tqdm()
- if index is None:
- raise ValueError("Failed to build an index for knn search")
- self._index = index
-
- logger.info("knn classifier index_infos", extra=self.index_infos)
+ d = train_embeddings.images.shape[-1]
+ self._index = IndexFlatL2(d)
+ self._index.add(train_embeddings.images)
 
  @staticmethod
  def get_default_params() -> dict[str, Any]:
@@ -65,7 +59,9 @@ def predict(self) -> tuple[ProbabilityArray, ClassArray]:
  # Calculate class votes from the distances (avoiding division by zero)
  # Note: Values stored in `dists` are the squared 2-norm values of the respective distance vectors
  max_value = np.finfo(np.float32).max
- scores = np.divide(1, dists, out=np.full_like(dists, max_value), where=dists != 0)
+ scores = np.divide(
+ 1, dists, out=np.full_like(dists, max_value), where=dists != 0
+ )
  # NOTE: if self.k and self.num_classes are both large, this might become a big one.
  # We can shape of a factor self.k if we count differently here.
  n = len(self._val_embeddings.images)