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

mini-nav/benchmarks/tasks/multi_object_retrieval.py

@@ -17,7 +17,7 @@ from configs.models import BenchmarkTaskConfig
 from rich.progress import track
 from torch import nn
 from torch.utils.data import DataLoader
-from transformers import BitImageProcessorFast
+from transformers import BitImageProcessor
 from utils.feature_extractor import extract_single_image_feature
 from utils.sam import load_sam_model, segment_image
 from utils.common import get_device
@@ -85,7 +85,7 @@ def _compute_scene_score(
         hit_rate = matched_count / len(query_object_ids)
 
         # Final score: sum_similarity * (hit_rate)^gamma
-        score = sum_similarity * (hit_rate ** gamma)
+        score = sum_similarity * (hit_rate**gamma)
         scene_scores[image_id] = score
 
     return scene_scores
@@ -143,7 +143,7 @@ class MultiObjectRetrievalTask(BaseBenchmarkTask):
     def build_database(
         self,
         model: nn.Module,
-        processor: BitImageProcessorFast,
+        processor: BitImageProcessor,
         train_dataset: Any,
         table: lancedb.table.Table,
         batch_size: int,
@@ -176,7 +176,9 @@ class MultiObjectRetrievalTask(BaseBenchmarkTask):
         record_id = 0
         records = []
 
-        for idx in track(range(len(train_dataset)), description="Building object database"):
+        for idx in track(
+            range(len(train_dataset)), description="Building object database"
+        ):
             item = train_dataset[idx]
             image = item["image"]
             image_id = item.get("image_id", f"image_{idx}")
@@ -204,13 +206,15 @@ class MultiObjectRetrievalTask(BaseBenchmarkTask):
                 object_id = f"{image_id}_obj_{mask_idx}"
                 category = mask_info.get("category", "unknown")
 
-                records.append({
-                    "id": record_id,
-                    "image_id": image_id,
-                    "object_id": object_id,
-                    "category": category,
-                    "vector": vector,
-                })
+                records.append(
+                    {
+                        "id": record_id,
+                        "image_id": image_id,
+                        "object_id": object_id,
+                        "category": category,
+                        "vector": vector,
+                    }
+                )
                 record_id += 1
 
         # Add all records to table
@@ -220,7 +224,7 @@ class MultiObjectRetrievalTask(BaseBenchmarkTask):
     def evaluate(
         self,
         model: nn.Module,
-        processor: BitImageProcessorFast,
+        processor: BitImageProcessor,
         test_dataset: Any,
         table: lancedb.table.Table,
         batch_size: int,
@@ -246,7 +250,9 @@ class MultiObjectRetrievalTask(BaseBenchmarkTask):
         correct = 0
         total = 0
 
-        for idx in track(range(len(test_dataset)), description=f"Evaluating Recall@{top_k}"):
+        for idx in track(
+            range(len(test_dataset)), description=f"Evaluating Recall@{top_k}"
+        ):
             item = test_dataset[idx]
             image = item["image"]
             target_image_id = item.get("image_id", f"image_{idx}")
@@ -295,7 +301,9 @@ class MultiObjectRetrievalTask(BaseBenchmarkTask):
                     retrieved_results[image_id].append((distance, object_id))
 
             # Compute scene scores
-            query_object_ids = [m.get("object_id", f"query_obj_{i}") for i, m in enumerate(query_masks)]
+            query_object_ids = [
+                m.get("object_id", f"query_obj_{i}") for i, m in enumerate(query_masks)
+            ]
             scene_scores = _compute_scene_score(
                 query_object_ids,
                 retrieved_results,
@@ -303,7 +311,9 @@ class MultiObjectRetrievalTask(BaseBenchmarkTask):
             )
 
             # Rank scenes by score
-            ranked_scenes = sorted(scene_scores.items(), key=lambda x: x[1], reverse=True)
+            ranked_scenes = sorted(
+                scene_scores.items(), key=lambda x: x[1], reverse=True
+            )
 
             # Check if target is in top-K
             top_k_scenes = [scene_id for scene_id, _ in ranked_scenes[:top_k]]
@@ -322,7 +332,7 @@ class MultiObjectRetrievalTask(BaseBenchmarkTask):
 
     def _infer_vector_dim(
         self,
-        processor: BitImageProcessorFast,
+        processor: BitImageProcessor,
         model: nn.Module,
         sample_image: Any,
     ) -> int:
@@ -347,7 +357,10 @@ class MultiObjectRetrievalTask(BaseBenchmarkTask):
         # Ensure mask is the right shape
         if mask.shape != image_np.shape[:2]:
             from skimage.transform import resize
-            mask_resized = resize(mask, image_np.shape[:2], order=0, anti_aliasing=False)
+
+            mask_resized = resize(
+                mask, image_np.shape[:2], order=0, anti_aliasing=False
+            )
         else:
             mask_resized = mask
 

mini-nav/benchmarks/tasks/retrieval.py

@@ -9,7 +9,7 @@ from benchmarks.tasks.registry import RegisterTask
 from rich.progress import track
 from torch import nn
 from torch.utils.data import DataLoader
-from transformers import BitImageProcessorFast
+from transformers import BitImageProcessor
 from utils.feature_extractor import extract_batch_features, infer_vector_dim
 
 
@@ -32,7 +32,7 @@ def _build_eval_schema(vector_dim: int) -> pa.Schema:
 
 
 def _establish_eval_database(
-    processor: BitImageProcessorFast,
+    processor: BitImageProcessor,
     model: nn.Module,
     table: lancedb.table.Table,
     dataloader: DataLoader,
@@ -69,7 +69,7 @@ def _establish_eval_database(
 
 
 def _evaluate_recall(
-    processor: BitImageProcessorFast,
+    processor: BitImageProcessor,
     model: nn.Module,
     table: lancedb.table.Table,
     dataloader: DataLoader,