misc: silence unnecessary progress bars

tti_eval/evaluation/image_retrieval.py

@@ -5,6 +5,7 @@
 from autofaiss import build_index
 
 from tti_eval.common import Embeddings
+from tti_eval.utils import disable_tqdm, enable_tqdm
 
 from .base import EvaluationModel
 
@@ -44,7 +45,9 @@ def __init__(
  self._class_counts = np.zeros(self.num_classes, dtype=np.int32)
  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

tti_eval/evaluation/knn.py

@@ -5,6 +5,7 @@
 from autofaiss import build_index
 
 from tti_eval.common import ClassArray, Embeddings, ProbabilityArray
+from tti_eval.utils import disable_tqdm, enable_tqdm
 
 from .base import ClassificationModel
 from .utils import softmax
@@ -42,13 +43,11 @@ 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,
+ 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

tti_eval/utils.py

@@ -1,10 +1,21 @@
+from functools import partialmethod
 from itertools import chain
 from typing import Literal, overload
 
+from tqdm import tqdm
+
 from tti_eval.common import EmbeddingDefinition
 from tti_eval.constants import PROJECT_PATHS
 
 
+def disable_tqdm():
+ tqdm.__init__ = partialmethod(tqdm.__init__, disable=True)
+
+
+def enable_tqdm():
+ tqdm.__init__ = partialmethod(tqdm.__init__, disable=False)
+
+
 @overload
 def read_all_cached_embeddings(as_list: Literal[True]) -> list[EmbeddingDefinition]:
  ...