SikongJueluo/Mini-Nav
1
0
Fork 0
mirror of https://github.com/SikongJueluo/Mini-Nav.git synced 2026-07-12 20:15:31 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
SikongJueluo
2026-04-11 15:20:38 +08:00
parent 01017277c3
commit 79b49f122a
8 changed files with 248 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
mini-nav/benchmarks/tasks/multi_object_retrieval.py
View File

@@ -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({
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

6
mini-nav/benchmarks/tasks/retrieval.py
View File

@@ -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,

8
mini-nav/commands/benchmark.py
View File

@@ -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

8
mini-nav/feature_retrieval.py
View File

@@ -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(

3
mini-nav/simulator/__init__.py
View File

@@ -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",

87
mini-nav/simulator/views.py
View File

@@ -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

12
mini-nav/utils/feature_extractor.py
View File

@@ -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)

110
notebooks/collect_test_images.py Normal file
View File

@@ -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()