feat(test): add collect test images notebook and replace BitImageProcessorFast

mini-nav/utils/feature_extractor.py

@@ -6,7 +6,7 @@ import torch
 from PIL import Image
 from torch import nn
 from torch.utils.data import DataLoader
-from transformers import BitImageProcessorFast
+from transformers import BitImageProcessor
 from rich.progress import track
 
 
@@ -27,7 +27,7 @@ def _extract_features_from_output(output: Any) -> torch.Tensor:
 
 
 def infer_vector_dim(
-    processor: BitImageProcessorFast,
+    processor: BitImageProcessor,
     model: nn.Module,
     sample_image: Any,
 ) -> int:
@@ -55,7 +55,7 @@ def infer_vector_dim(
 
 @torch.no_grad()
 def extract_single_image_feature(
-    processor: BitImageProcessorFast,
+    processor: BitImageProcessor,
     model: nn.Module,
     image: Union[Image.Image, Any],
 ) -> List[float]:
@@ -82,7 +82,7 @@ def extract_single_image_feature(
 
 @torch.no_grad()
 def extract_batch_features(
-    processor: BitImageProcessorFast,
+    processor: BitImageProcessor,
     model: nn.Module,
     images: Union[List[Any], Any],
     batch_size: int = 32,
@@ -115,7 +115,9 @@ def extract_batch_features(
 
     # Handle list of images
     all_features = []
-    for i in track(range(0, len(images), batch_size), description="Extracting features"):
+    for i in track(
+        range(0, len(images), batch_size), description="Extracting features"
+    ):
         batch_imgs = images[i : i + batch_size]
         inputs = processor(images=batch_imgs, return_tensors="pt")
         inputs.to(device)