diff --git a/mini-nav/benchmarks/tasks/multi_object_retrieval.py b/mini-nav/benchmarks/tasks/multi_object_retrieval.py index 75fafe5..3758acb 100644 --- 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 diff --git a/mini-nav/benchmarks/tasks/retrieval.py b/mini-nav/benchmarks/tasks/retrieval.py index 435768a..c63e41a 100644 --- 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, diff --git a/mini-nav/commands/benchmark.py b/mini-nav/commands/benchmark.py index 850bf73..20809f5 100644 --- 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 diff --git a/mini-nav/feature_retrieval.py b/mini-nav/feature_retrieval.py index 1e706c6..f2c8bea 100644 --- 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( diff --git a/mini-nav/simulator/__init__.py b/mini-nav/simulator/__init__.py index a5e55e1..170f74f 100644 --- 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", diff --git a/mini-nav/simulator/views.py b/mini-nav/simulator/views.py index 47c4be2..d61b13b 100644 --- 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 diff --git a/mini-nav/utils/feature_extractor.py b/mini-nav/utils/feature_extractor.py index bbed324..a17dbc9 100644 --- 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) diff --git a/notebooks/collect_test_images.py b/notebooks/collect_test_images.py new file mode 100644 index 0000000..15dcb1c --- /dev/null +++ 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()