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

47
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({
"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

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)