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

2026-07-12 20:15:31 +08:00 · 2026-04-11 15:20:38 +08:00
parent 01017277c3
commit 79b49f122a
8 changed files with 248 additions and 33 deletions
--- a/mini-nav/benchmarks/tasks/multi_object_retrieval.py
+++ b/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

--- a/mini-nav/benchmarks/tasks/retrieval.py
+++ b/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,
--- a/mini-nav/commands/benchmark.py
+++ b/mini-nav/commands/benchmark.py
@@ -15,7 +15,7 @@ def benchmark(
    import torch.nn.functional as F
    from benchmarks import run_benchmark
    from configs import cfg_manager
-    from transformers import AutoImageProcessor, AutoModel, BitImageProcessorFast
+    from transformers import AutoImageProcessor, AutoModel, BitImageProcessor
    from utils import get_device

    config = cfg_manager.get()
@@ -25,7 +25,7 @@ def benchmark(

    model_cfg = config.model
    processor = cast(
-        BitImageProcessorFast,
+        BitImageProcessor,
        AutoImageProcessor.from_pretrained(model_cfg.dino_model, device_map=device),
    )

@@ -63,7 +63,9 @@ def benchmark(
        def encode(self, images: list) -> torch.Tensor:
            if self.compressor is None:
                return self.extract_features(images)
-            tokens = self.dino(**processor(images, return_tensors="pt").to(device)).last_hidden_state
+            tokens = self.dino(
+                **processor(images, return_tensors="pt").to(device)
+            ).last_hidden_state
            _, _, bits = self.compressor(tokens)
            return bits

--- a/mini-nav/feature_retrieval.py
+++ b/mini-nav/feature_retrieval.py
@@ -11,7 +11,7 @@ from torch import nn
 from transformers import (
    AutoImageProcessor,
    AutoModel,
-    BitImageProcessorFast,
+    BitImageProcessor,
    Dinov2Model,
 )
 from utils.feature_extractor import extract_batch_features
@@ -38,7 +38,7 @@ class FeatureRetrieval:
    _instance: Optional["FeatureRetrieval"] = None

    _initialized: bool = False
-    processor: BitImageProcessorFast
+    processor: BitImageProcessor
    model: nn.Module

    def __new__(cls, *args, **kwargs) -> "FeatureRetrieval":
@@ -48,7 +48,7 @@ class FeatureRetrieval:

    def __init__(
        self,
-        processor: Optional[BitImageProcessorFast] = None,
+        processor: Optional[BitImageProcessor] = None,
        model: Optional[nn.Module] = None,
    ) -> None:
        """Initialize the singleton with processor and model.
@@ -124,7 +124,7 @@ if __name__ == "__main__":
    ]

    processor = cast(
-        BitImageProcessorFast,
+        BitImageProcessor,
        AutoImageProcessor.from_pretrained("facebook/dinov2-large", device_map="cuda"),
    )
    model = cast(
--- a/mini-nav/simulator/init.py
+++ b/mini-nav/simulator/init.py
@@ -5,7 +5,7 @@ from .habitat import (
 )
 from .image_save import save_object_image, save_room_view
 from .topdown import TopDownRenderStyle, TopDownSceneElements, render_topdown_scene_map
-from .views import RoomViewsByRoom, collect_room_views_by_room
+from .views import RoomViewsByRoom, collect_room_views_by_room, collect_scene_images

 __all__ = [
    "HabitatSimulatorConfig",
@@ -14,6 +14,7 @@ __all__ = [
    "TopDownSceneElements",
    "close_habitat_simulator",
    "collect_room_views_by_room",
+    "collect_scene_images",
    "create_habitat_simulator",
    "render_topdown_scene_map",
    "save_object_image",
--- a/mini-nav/simulator/views.py
+++ b/mini-nav/simulator/views.py
@@ -1,8 +1,10 @@
 from __future__ import annotations

 from importlib import import_module
+from pathlib import Path
 from typing import Any, Callable, Iterable, Sequence

+import numpy as np
 from rich.progress import track

 RoomViewsByRoom = dict[str, list[Any]]
@@ -42,3 +44,88 @@ def collect_room_views_by_room(
        all_room_views[room_node.room_id] = room_views

    return all_room_views
+
+
+def collect_scene_images(
+    scene_name: str,
+    scene_path: str,
+    output_dir: Path,
+    *,
+    image_size: int = 1024,
+    views_per_point: int = 12,
+    points_per_scene: int = 5,
+    seed: int = 42,
+) -> int:
+    """Collect RGB images from random navigable points in a scene.
+
+    Creates a Habitat simulator for the given scene, samples random
+    navigable points, and captures rotated views at each point. Images
+    are saved as PNG files under ``output_dir / scene_name / {point:03d} /``.
+
+    Args:
+        scene_name: Identifier used as subdirectory name.
+        scene_path: Path to the Habitat scene dataset file (.glb).
+        output_dir: Root output directory for saved images.
+        image_size: Resolution (width and height) of captured images.
+        views_per_point: Number of views captured at each point.
+        points_per_scene: Number of random points to sample.
+        seed: Seed for pathfinder reproducibility.
+
+    Returns:
+        Number of images successfully collected.
+    """
+    from utils.image import numpy_to_pil
+
+    from .habitat import HabitatSimulatorConfig, create_habitat_simulator
+
+    config = HabitatSimulatorConfig(
+        scene_path=scene_path,
+        image_size=image_size,
+        views_per_room=views_per_point,
+    )
+    sim, agent = create_habitat_simulator(config)
+    sim.pathfinder.seed(seed)
+
+    collected_count = 0
+    try:
+        for point_idx in range(points_per_scene):
+            point = None
+            for _ in range(10):
+                candidate = sim.pathfinder.get_random_navigable_point()
+                candidate = np.asarray(candidate, dtype=np.float32)
+                if not np.isfinite(candidate).all():
+                    continue
+                if not sim.pathfinder.is_navigable(candidate):
+                    continue
+                point = candidate
+                break
+
+            if point is None:
+                print(
+                    f"[WARN] Skip {scene_name} point {point_idx:03d}: "
+                    "no valid navigable point"
+                )
+                continue
+
+            agent_state = agent.get_state()
+            agent_state.position = point
+            agent.set_state(agent_state)
+
+            for view_idx in range(views_per_point):
+                obs = sim.get_sensor_observations()
+                rgb = obs["color_sensor"]
+                image = numpy_to_pil(rgb)
+                save_path = (
+                    output_dir
+                    / scene_name
+                    / f"{point_idx:03d}"
+                    / f"view_{view_idx:03d}.png"
+                )
+                save_path.parent.mkdir(parents=True, exist_ok=True)
+                image.save(str(save_path))
+                collected_count += 1
+                sim.step("turn_left")
+    finally:
+        sim.close()
+
+    return collected_count
--- a/mini-nav/utils/feature_extractor.py
+++ b/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)
--- a/notebooks/collect_test_images.py
+++ b/notebooks/collect_test_images.py
@@ -0,0 +1,110 @@
+# /// script
+# requires-python = ">=3.13"
+# dependencies = [
+#     "marimo>=0.21.1",
+#     "pyzmq>=27.1.0",
+# ]
+# ///
+
+# pyright: reportMissingImports=false, reportMissingParameterType=false, reportUnknownParameterType=false, reportUnknownVariableType=false, reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnusedParameter=false, reportUnusedCallResult=false, reportUnusedExpression=false
+
+import marimo
+
+__generated_with = "0.22.4"
+app = marimo.App(width="medium", app_title="Test Scene Image Collector")
+
+
+@app.cell
+def _():
+    from pathlib import Path
+
+    import marimo as mo
+
+    return Path, mo
+
+
+@app.cell
+def _(Path, mo):
+    SCENES = {
+        "skokloster-castle": "data/scene_datasets/habitat-test-scenes/skokloster-castle.glb",
+        "apartment_1": "data/scene_datasets/habitat-test-scenes/apartment_1.glb",
+        "van-gogh-room": "data/scene_datasets/habitat-test-scenes/van-gogh-room.glb",
+    }
+    IMAGE_SIZE = 1024
+    VIEWS_PER_POINT = 12
+    POINTS_PER_SCENE = 5
+    SEED = 42
+    OUTPUT_DIR = Path("outputs/test_image")
+
+    mo.md(
+        f"## Configuration\n- Scenes: {len(SCENES)}\n- Points/scene: {POINTS_PER_SCENE}\n- Views/point: {VIEWS_PER_POINT}\n- Resolution: {IMAGE_SIZE}x{IMAGE_SIZE}\n- Total images: {len(SCENES) * POINTS_PER_SCENE * VIEWS_PER_POINT}"
+    )
+    return (
+        IMAGE_SIZE,
+        OUTPUT_DIR,
+        POINTS_PER_SCENE,
+        SCENES,
+        SEED,
+        VIEWS_PER_POINT,
+    )
+
+
+@app.cell
+def _(IMAGE_SIZE, OUTPUT_DIR, POINTS_PER_SCENE, SCENES, SEED, VIEWS_PER_POINT):
+    from simulator import collect_scene_images
+
+    scene_stats = {}
+    total_images = 0
+
+    for scene_name, scene_path in SCENES.items():
+        print(f"Collecting {scene_name}...")
+        collected = collect_scene_images(
+            scene_name=scene_name,
+            scene_path=scene_path,
+            output_dir=OUTPUT_DIR,
+            image_size=IMAGE_SIZE,
+            views_per_point=VIEWS_PER_POINT,
+            points_per_scene=POINTS_PER_SCENE,
+            seed=SEED,
+        )
+        scene_stats[scene_name] = collected
+        total_images += collected
+        print(f"Collected {collected} images for {scene_name}")
+    return scene_stats, total_images
+
+
+@app.cell
+def _(OUTPUT_DIR, mo, scene_stats, total_images):
+    mo.md(f"""
+    ## Summary\n- Output directory: `{OUTPUT_DIR}`\n- Total images: {total_images}\n- Per-scene: {scene_stats}
+    """)
+    return
+
+
+@app.cell
+def _(OUTPUT_DIR, SCENES, mo):
+    from PIL import Image
+
+    preview_grid = []
+    for _scene_name in SCENES:
+        sample_path = OUTPUT_DIR / _scene_name / "000" / "view_000.png"
+        if not sample_path.exists():
+            continue
+        preview_grid.append(
+            mo.vstack(
+                [mo.md(f"**{_scene_name}**"), mo.image(Image.open(sample_path))],
+                align="center",
+            )
+        )
+
+    preview = (
+        mo.vstack(preview_grid, align="center")
+        if preview_grid
+        else mo.md("No preview images yet.")
+    )
+    preview
+    return
+
+
+if __name__ == "__main__":
+    app.run()