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(verification): add batch segmentation and image saving

Browse Source
This commit is contained in:
SikongJueluo
2026-03-28 21:30:02 +08:00
parent f604c85a79
commit f6c1a67e88
4 changed files with 182 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
mini-nav/simulator/habitat.py
View File

@@ -9,7 +9,7 @@ from typing import Any
class HabitatSimulatorConfig:
scene_path: str
views_per_room: int = 6
image_size: int = 256
image_size: int = 512
sensor_height: float = 1.5
move_forward_step: float = 0.25
enable_physics: bool = False

57
mini-nav/tests/test_image_utils.py
View File

@@ -83,19 +83,17 @@ def test_segment_image_filters_tensor_masks_by_min_area() -> None:
def test_segment_image_dataset_returns_per_image_masks_in_order() -> None:
first_masks = {
"masks": torch.tensor(
first_masks = torch.tensor(
[[[1, 1, 0], [1, 1, 0], [0, 0, 0]]],
dtype=torch.float32,
)
}
second_masks = {
"masks": torch.tensor(
second_masks = torch.tensor(
[[[1, 1, 1], [1, 1, 1], [1, 1, 1]]],
dtype=torch.float32,
)
}
mock_generator = Mock(side_effect=[first_masks, second_masks])
mock_generator = Mock(
return_value=[{"masks": first_masks}, {"masks": second_masks}]
)
images = [
Image.new("RGB", (3, 3), color=(0, 0, 0)),
Image.new("RGB", (3, 3), color=(0, 0, 0)),
@@ -112,4 +110,47 @@ def test_segment_image_dataset_returns_per_image_masks_in_order() -> None:
assert len(result) == 2
assert result[0][0]["area"] == 4
assert result[1][0]["area"] == 9
assert mock_generator.call_count == 2
assert mock_generator.call_count == 1
def test_segment_image_dataset_falls_back_to_single_image_calls() -> None:
call_index = {"value": 0}
def fake_generator(images, points_per_batch):
if isinstance(images, list):
raise TypeError("Batch input unsupported")
result_options = [
{
"masks": torch.tensor(
[[[1, 1, 0], [1, 1, 0], [0, 0, 0]]],
dtype=torch.float32,
)
},
{
"masks": torch.tensor(
[[[1, 1, 1], [1, 1, 1], [1, 1, 1]]],
dtype=torch.float32,
)
},
]
out = result_options[call_index["value"]]
call_index["value"] += 1
return out
images = [
Image.new("RGB", (3, 3), color=(0, 0, 0)),
Image.new("RGB", (3, 3), color=(0, 0, 0)),
]
result = segment_image_dataset(
fake_generator,
images,
min_area=2,
max_masks=5,
points_per_batch=16,
)
assert len(result) == 2
assert result[0][0]["area"] == 4
assert result[1][0]["area"] == 9

98
mini-nav/utils/image.py
View File

@@ -29,10 +29,79 @@ def segment_image(
"""
image_rgb = image.convert("RGB")
raw_output = mask_generator(image_rgb, points_per_batch=points_per_batch)
return _normalize_and_filter_masks(
raw_output, min_area=min_area, max_masks=max_masks
)
def segment_image_dataset(
mask_generator: Any,
images: Sequence[Image.Image],
min_area: int = 32 * 32,
max_masks: int = 5,
points_per_batch: int = 64,
) -> list[list[dict[str, Any]]]:
image_list = list(images)
if not image_list:
return []
image_rgb_list = [image.convert("RGB") for image in image_list]
try:
raw_batch_output = mask_generator(
image_rgb_list,
points_per_batch=points_per_batch,
)
batch_items = _split_batch_output(
raw_batch_output, expected_size=len(image_list)
)
if batch_items is not None:
return [
_normalize_and_filter_masks(
batch_item,
min_area=min_area,
max_masks=max_masks,
)
for batch_item in batch_items
]
except TypeError:
pass
return [
_normalize_and_filter_masks(
mask_generator(image_rgb, points_per_batch=points_per_batch),
min_area=min_area,
max_masks=max_masks,
)
for image_rgb in image_rgb_list
]
def _split_batch_output(raw_output: Any, expected_size: int) -> list[Any] | None:
if isinstance(raw_output, list):
if len(raw_output) == expected_size:
return raw_output
return None
if isinstance(raw_output, dict):
raw_masks = raw_output.get("masks", raw_output)
if isinstance(raw_masks, list) and len(raw_masks) == expected_size:
return raw_masks
return None
def _normalize_and_filter_masks(
raw_output: Any,
min_area: int,
max_masks: int,
) -> list[dict[str, Any]]:
raw_masks = (
raw_output.get("masks", raw_output)
if isinstance(raw_output, dict)
else raw_output
)
normalized_masks: list[dict[str, Any]] = []
if isinstance(raw_masks, list):
if raw_masks and isinstance(raw_masks[0], dict):
normalized_masks = raw_masks
@@ -55,35 +124,16 @@ def segment_image(
if not normalized_masks:
return []
filtered_masks = [m for m in normalized_masks if int(m["area"]) >= min_area]
filtered_masks = [
mask for mask in normalized_masks if int(mask["area"]) >= min_area
]
if not filtered_masks:
return []
sorted_masks = sorted(filtered_masks, key=lambda x: x["area"], reverse=True)
sorted_masks = sorted(filtered_masks, key=lambda mask: mask["area"], reverse=True)
return sorted_masks[:max_masks]
def segment_image_dataset(
mask_generator: Any,
images: Sequence[Image.Image],
min_area: int = 32 * 32,
max_masks: int = 5,
points_per_batch: int = 64,
) -> list[list[dict[str, Any]]]:
image_list = list(images)
return [
segment_image(
mask_generator,
image,
min_area=min_area,
max_masks=max_masks,
points_per_batch=points_per_batch,
)
for image in image_list
]
def _to_numpy_mask_array(mask_like: Any) -> np.ndarray | None:
if mask_like is None:
return None

69
notebooks/verification.py
View File

@@ -23,6 +23,7 @@ def import_packages():
from matplotlib import pyplot as plt
from PIL import Image
from configs import cfg_manager
from compressors.pipeline import HashPipeline
from scenegraph import ObjectNode, RoomNode, SimpleSceneGraph
from simulator import (
@@ -32,7 +33,7 @@ def import_packages():
create_habitat_simulator,
render_topdown_scene_map,
)
from utils.image import extract_masked_region, segment_image
from utils.image import extract_masked_region, segment_image_dataset
return (
HabitatSimulatorConfig,
@@ -44,6 +45,7 @@ def import_packages():
TopDownSceneElements,
collect_room_views_by_room,
create_habitat_simulator,
cfg_manager,
extract_masked_region,
maps,
mo,
@@ -51,7 +53,7 @@ def import_packages():
pl,
plt,
render_topdown_scene_map,
segment_image,
segment_image_dataset,
)
@@ -139,6 +141,7 @@ def build_scene_graph_pipeline(
SimpleSceneGraph,
agent,
collect_room_views_by_room,
cfg_manager,
extract_masked_region,
hash_bits,
mo,
@@ -147,7 +150,7 @@ def build_scene_graph_pipeline(
room_nodes,
sam_max_masks,
sam_min_area,
segment_image,
segment_image_dataset,
sim,
views_per_room,
):
@@ -170,6 +173,9 @@ def build_scene_graph_pipeline(
rooms={room.room_id: room for room in room_nodes},
objects={},
)
verification_output_dir = cfg_manager.get().output.directory / "verification"
verification_output_dir.mkdir(parents=True, exist_ok=True)
total_masks = 0
object_index = 0
@@ -180,31 +186,48 @@ def build_scene_graph_pipeline(
]
object_dataset = []
for room_id, _view_idx, rgb in mo.status.progress_bar(
room_view_dataset,
room_view_images = []
for _, _, rgb in room_view_dataset:
rgb3 = rgb[..., :3] if rgb.shape[-1] > 3 else rgb
room_view_images.append(Image.fromarray(rgb3.astype(np.uint8)))
masks_dataset = segment_image_dataset(
hash_pipeline.mask_generator,
room_view_images,
min_area=hash_pipeline.sam_min_mask_area,
max_masks=hash_pipeline.sam_max_masks,
points_per_batch=hash_pipeline.sam_points_per_batch,
)
if len(masks_dataset) != len(room_view_dataset):
raise RuntimeError("SAM dataset output size mismatch with room_view_dataset.")
dataset_jobs = list(zip(room_view_dataset, room_view_images, masks_dataset))
for (room_id, view_idx, _), image, masks in mo.status.progress_bar(
dataset_jobs,
title="Building object dataset",
subtitle="Running SAM segmentation",
show_eta=True,
show_rate=True,
):
rgb3 = rgb[..., :3] if rgb.shape[-1] > 3 else rgb
image = Image.fromarray(rgb3.astype(np.uint8))
room_output_dir = verification_output_dir / room_id
room_output_dir.mkdir(parents=True, exist_ok=True)
room_view_path = room_output_dir / f"view_{view_idx:03d}.png"
image.convert("RGB").save(room_view_path, format="PNG")
masks = segment_image(
hash_pipeline.mask_generator,
image,
min_area=hash_pipeline.sam_min_mask_area,
max_masks=hash_pipeline.sam_max_masks,
points_per_batch=hash_pipeline.sam_points_per_batch,
)
total_masks += len(masks)
for mask in masks:
for mask_idx, mask in enumerate(masks):
masked_image = extract_masked_region(image, mask["segment"])
object_dataset.append((room_id, mask["bbox"], masked_image))
object_dataset.append(
(room_id, view_idx, mask_idx, mask["bbox"], masked_image)
)
if object_dataset:
masked_images = [item[2] for item in object_dataset]
masked_images = [item[4] for item in object_dataset]
if any(not isinstance(img, Image.Image) for img in masked_images):
raise TypeError(
"object_dataset contains non-image entries for batch inference."
)
batched_bits = hash_pipeline.forward_dataset(
masked_images,
batch_size=pipeline_batch_size,
@@ -217,7 +240,9 @@ def build_scene_graph_pipeline(
else:
batched_bits = []
for ob_idx, (room_id, bbox, _) in enumerate(object_dataset):
for ob_idx, (room_id, view_idx, mask_idx, bbox, masked_image) in enumerate(
object_dataset
):
bits = batched_bits[ob_idx]
obj_center = np.array(
[bbox[0] + bbox[2] / 2, bbox[1] + bbox[3] / 2, 0.0],
@@ -227,6 +252,13 @@ def build_scene_graph_pipeline(
obj_id = f"obj_{object_index:04d}"
object_index += 1
room_output_dir = verification_output_dir / room_id
room_output_dir.mkdir(parents=True, exist_ok=True)
object_image_path = (
room_output_dir / f"{obj_id}_view{view_idx:03d}_mask{mask_idx:02d}.png"
)
masked_image.convert("RGB").save(object_image_path, format="PNG")
bits_array = np.asarray(bits.detach().cpu().numpy()).reshape(-1)
if bits_array.size == 512:
bits_binary = (bits_array > 0).astype(np.uint8)
@@ -250,6 +282,7 @@ def build_scene_graph_pipeline(
print(f"Total objects created: {len(scene_graph.objects)}")
print(f"Total processed masks: {total_masks}")
print(f"Saved object images to: {verification_output_dir}")
return (scene_graph,)