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(notebooks): add proposal segmentation notebook and enhance verification

Browse Source
This commit is contained in:
SikongJueluo
2026-04-01 11:53:51 +08:00
parent f06b0625b4
commit aedcb25610
3 changed files with 415 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

383
notebooks/proposal_segament.py Normal file
View File

@@ -0,0 +1,383 @@
import marimo
__generated_with = "0.21.1"
app = marimo.App()
@app.cell
def _():
import marimo as mo
import numpy as np
import polars as pl
from PIL import Image, ImageDraw, ImageFont
import torch
import io
import requests
return Image, ImageDraw, ImageFont, mo, np, torch
@app.cell
def _(np):
from scenegraph import RoomNode
from simulator import (
HabitatSimulatorConfig,
TopDownSceneElements,
create_habitat_simulator,
render_topdown_scene_map,
)
from habitat.utils.visualizations import maps
scene_path = "data/scene_datasets/habitat-test-scenes/skokloster-castle.glb"
image_size = 768
num_rooms = 3
views_per_room = 6
meters_per_pixel = 0.05
sim, agent = create_habitat_simulator(
HabitatSimulatorConfig(
scene_path=scene_path,
views_per_room=views_per_room,
image_size=image_size,
sensor_height=1.5,
move_forward_step=0.25,
enable_physics=False,
)
)
room_nodes = []
for _idx in range(num_rooms):
point = sim.pathfinder.get_random_navigable_point()
room_nodes.append(
RoomNode(
room_id=f"room_{_idx:02d}",
center=np.asarray(point, dtype=np.float32),
bbox_extent=np.asarray([1.5, 2.0, 1.5], dtype=np.float32),
)
)
render_topdown_scene_map(
pathfinder=sim.pathfinder,
elements=TopDownSceneElements(room_nodes=room_nodes),
meters_per_pixel=meters_per_pixel,
maps_module=maps,
)
return agent, room_nodes, sim, views_per_room
@app.cell
def _(agent, room_nodes, sim, views_per_room):
from rich.progress import track
from compressors import MaskScoringConfig, rank_masks
from compressors.pipeline import HashPipeline
from compressors.proposal import (
extract_masked_region,
generate_proposals_batch,
)
from simulator import save_object_image, save_room_view
from simulator import collect_room_views_by_room
all_room_views = collect_room_views_by_room(
agent=agent,
sim=sim,
room_nodes=room_nodes,
views_per_room=views_per_room,
)
return (all_room_views,)
@app.cell
def _(ImageDraw, ImageFont, all_room_views, mo, torch):
from transformers import Owlv2Processor, Owlv2ForObjectDetection
from utils import get_device, numpy_to_pil
device = get_device()
# 你真正想控制的显示阈值
score_threshold = 0.25
# 1. 加载模型
processor = Owlv2Processor.from_pretrained(
"google/owlv2-base-patch16-ensemble"
)
model = Owlv2ForObjectDetection.from_pretrained(
"google/owlv2-base-patch16-ensemble"
).to(device)
model.eval()
# 2. 读取图片
image = numpy_to_pil(all_room_views["room_00"][1])
# 3. 文本查询
text_labels = [
[
"a cat",
"a dog",
"a TV remote control",
"a chair",
"a table",
"a vase",
"a painting",
"a window"
]
]
# 4. 推理
inputs = processor(text=text_labels, images=image, return_tensors="pt").to(
device
)
with torch.no_grad():
outputs = model(**inputs)
# 5. 后处理
target_sizes = torch.tensor([(image.height, image.width)])
results = processor.post_process_grounded_object_detection(
outputs=outputs,
target_sizes=target_sizes,
threshold=0.1, # 模型后处理的第一层粗筛
text_labels=text_labels,
)
result = results[0]
boxes = result["boxes"]
scores = result["scores"]
pred_labels = result["text_labels"]
# 6. 二次过滤：只保留高分目标
filtered_items = [
(_box, _score, _text_label)
for _box, _score, _text_label in zip(boxes, scores, pred_labels)
if _score.item() >= score_threshold
]
# 7. 画框
_vis_image = image.copy()
_draw = ImageDraw.Draw(_vis_image)
try:
font = ImageFont.truetype("arial.ttf", 18)
except Exception:
font = ImageFont.load_default()
for _box, _score, _text_label in filtered_items:
_x1, _y1, _x2, _y2 = [float(v) for v in _box.tolist()]
_label = f"{_text_label}: {_score.item():.3f}"
# 边框
_draw.rectangle([_x1, _y1, _x2, _y2], outline="red", width=3)
# 文本背景框
try:
_tx1, _ty1, _tx2, _ty2 = _draw.textbbox((_x1, _y1), _label, font=font)
except Exception:
_w = _draw.textlength(_label, font=font)
_h = 20
_tx1, _ty1, _tx2, _ty2 = _x1, _y1, _x1 + _w + 6, _y1 + _h
text_bg = [_tx1, max(0, _ty1 - 2), _tx2 + 4, _ty2 + 2]
_draw.rectangle(text_bg, fill="red")
_draw.text((_x1 + 2, max(0, _y1)), _label, fill="white", font=font)
# 8. 结果文本
detection_lines = []
for _box, _score, _text_label in filtered_items:
box_rounded = [round(_v, 2) for _v in _box.tolist()]
detection_lines.append(
f"- {_text_label}: score={_score.item():.3f}, box={box_rounded}"
)
if not detection_lines:
detection_text = f"没有检测到 score >= {score_threshold:.2f} 的目标"
else:
detection_text = "\n".join(detection_lines)
# 9. marimo 输出
mo.vstack(
[
mo.md(
f"## OWLv2 检测可视化结果\n\n过滤阈值：`score >= {score_threshold:.2f}`"
),
mo.image(_vis_image, width=700),
mo.md(detection_text),
]
)
return device, filtered_items, image
@app.cell
def _(Image, ImageDraw, device, filtered_items, image, mo, np, torch):
from transformers import Sam2Processor, Sam2Model
# -----------------------------
# 0. 加载 SAM2
# -----------------------------
sam2_model_id = "facebook/sam2.1-hiera-large"
sam2_processor = Sam2Processor.from_pretrained(sam2_model_id)
sam2_model = Sam2Model.from_pretrained(sam2_model_id).to(device)
sam2_model.eval()
# -----------------------------
# 1. 从上一格的 OWLv2 结果中取框
# filtered_items: [(box, score, text_label), ...]
# SAM2 需要的 box 格式：
# [[[x1, y1, x2, y2], [x1, y1, x2, y2], ...]]
# 这里最外层是 batch 维度，因为我们当前只处理 1 张图
# -----------------------------
input_boxes = [
[
[float(v) for v in box.tolist()]
for box, score, text_label in filtered_items
]
]
# 没有框就直接提前返回
if len(input_boxes[0]) == 0:
mo.vstack(
[
mo.md("## SAM2 分割结果"),
mo.md(
"没有可用于分割的检测框，请先降低 OWLv2 的 score_threshold 或检查检测结果。"
),
]
)
else:
# -----------------------------
# 2. 预处理并推理
# 官方文档支持 bounding box prompt
# 并且 multimask_output=False 时，每个框只返回一张最佳 mask
# -----------------------------
sam2_inputs = sam2_processor(
images=image,
input_boxes=input_boxes,
return_tensors="pt",
).to(device)
with torch.no_grad():
sam2_outputs = sam2_model(**sam2_inputs, multimask_output=False)
# -----------------------------
# 3. 后处理到原图尺寸
# 文档说明 pred_masks 需要 post_process_masks
# -----------------------------
all_masks = sam2_processor.post_process_masks(
sam2_outputs.pred_masks.cpu(),
sam2_inputs["original_sizes"],
)
# 当前只有 1 张图，所以取第 0 个
# 形状通常是 [num_objects, num_masks, H, W]
# 由于 multimask_output=False，num_masks 通常为 1
masks_for_image = all_masks[0]
# 兼容处理：压掉单 mask 维度
# 目标形状变成 [num_objects, H, W]
if masks_for_image.ndim == 4 and masks_for_image.shape[1] == 1:
masks_for_image = masks_for_image[:, 0]
# IoU 分数，文档中也给出了 outputs.iou_scores
# 形状通常是 [batch_size, point_batch_size, num_masks]
iou_scores = sam2_outputs.iou_scores.detach().cpu()[0]
if iou_scores.ndim == 2 and iou_scores.shape[1] == 1:
iou_scores = iou_scores[:, 0]
# -----------------------------
# 4. 可视化：把 mask 半透明叠加到原图上
# -----------------------------
base_rgba = image.convert("RGBA")
overlay = Image.new("RGBA", image.size, (0, 0, 0, 0))
overlay_draw = ImageDraw.Draw(overlay)
# 准备一个可重复的颜色列表
colors = [
(255, 0, 0, 90),
(0, 255, 0, 90),
(0, 0, 255, 90),
(255, 255, 0, 90),
(255, 0, 255, 90),
(0, 255, 255, 90),
(255, 128, 0, 90),
(128, 0, 255, 90),
]
# 先画 mask，再画 box
_vis_image = base_rgba.copy()
for _idx, ((_box, _score, _text_label), mask) in enumerate(
zip(filtered_items, masks_for_image)
):
color = colors[_idx % len(colors)]
# mask -> numpy bool
mask_np = mask.numpy() > 0
# 做一层彩色 mask
mask_rgba = np.zeros((image.height, image.width, 4), dtype=np.uint8)
mask_rgba[mask_np] = color
mask_img = Image.fromarray(mask_rgba, mode="RGBA")
_vis_image = Image.alpha_composite(_vis_image, mask_img)
# 画框和标签
_draw = ImageDraw.Draw(_vis_image)
for _idx, ((_box, _score, _text_label), mask) in enumerate(
zip(filtered_items, masks_for_image)
):
_x1, _y1, _x2, _y2 = [float(v) for v in _box.tolist()]
iou_val = (
float(iou_scores[_idx].item())
if _idx < len(iou_scores)
else float("nan")
)
_label = (
f"{_text_label} | owl={_score.item():.3f} | sam2_iou={iou_val:.3f}"
)
_draw.rectangle([_x1, _y1, _x2, _y2], outline=(255, 0, 0, 255), width=3)
try:
_tx1, _ty1, _tx2, _ty2 = _draw.textbbox((_x1, _y1), _label)
except Exception:
_tx1, _ty1, _tx2, _ty2 = _x1, _y1, _x1 + 220, _y1 + 20
_draw.rectangle(
[_tx1, max(0, _ty1 - 2), _tx2 + 4, _ty2 + 2],
fill=(255, 0, 0, 220),
)
_draw.text((_x1 + 2, max(0, _y1)), _label, fill="white")
# -----------------------------
# 5. 输出摘要
# -----------------------------
summary_lines = []
for _idx, ((_box, _score, _text_label), mask) in enumerate(
zip(filtered_items, masks_for_image)
):
mask_area = int((mask.numpy() > 0).sum())
iou_val = (
float(iou_scores[_idx].item())
if _idx < len(iou_scores)
else float("nan")
)
summary_lines.append(
f"- {_text_label}: owl_score={_score.item():.3f}, sam2_iou={iou_val:.3f}, mask_area={mask_area}"
)
mo.vstack(
[
mo.md("## SAM2 分割可视化结果"),
mo.image(_vis_image, width=700),
mo.md(
"\n".join(summary_lines) if summary_lines else "没有生成任何 mask"
),
]
)
return
if __name__ == "__main__":
app.run()

33
notebooks/verification.py
View File

@@ -61,12 +61,13 @@ def setup_verification_context(
np,
):
scene_path = "data/scene_datasets/habitat-test-scenes/skokloster-castle.glb"
image_size = 256
image_size = 512
num_rooms = 4
views_per_room = 6
meters_per_pixel = 0.05
sam_max_masks = 5
sam_candidate_masks = 24
sam_min_area = 32 * 32
hash_bits = 512
pipeline_batch_size = 64
@@ -102,6 +103,7 @@ def setup_verification_context(
meters_per_pixel,
pipeline_batch_size,
room_nodes,
sam_candidate_masks,
sam_max_masks,
sam_min_area,
sim,
@@ -142,6 +144,7 @@ def build_scene_graph_pipeline(
np,
pipeline_batch_size,
room_nodes,
sam_candidate_masks,
sam_max_masks,
sam_min_area,
sim,
@@ -149,6 +152,7 @@ def build_scene_graph_pipeline(
):
from rich.progress import track
from compressors import MaskScoringConfig, rank_masks
from compressors.pipeline import HashPipeline
from compressors.proposal import extract_masked_region, generate_proposals_batch
from simulator import save_object_image, save_room_view
@@ -176,7 +180,17 @@ def build_scene_graph_pipeline(
verification_output_dir = cfg_manager.get().output.directory / "verification"
verification_output_dir.mkdir(parents=True, exist_ok=True)
mask_scoring_config = MaskScoringConfig(
max_area_ratio=0.45,
reject_edge_touch_count=3,
reject_large_edge_touch_count=2,
reject_large_edge_area_ratio=0.12,
max_components=4,
min_largest_component_ratio=0.70,
)
total_masks = 0
total_raw_masks = 0
object_index = 0
room_view_dataset = [
@@ -192,7 +206,7 @@ def build_scene_graph_pipeline(
hash_pipeline.mask_generator,
room_view_images,
min_area=hash_pipeline.sam_min_mask_area,
max_masks=hash_pipeline.sam_max_masks,
max_masks=sam_candidate_masks,
points_per_batch=hash_pipeline.sam_points_per_batch,
)
if len(masks_dataset) != len(room_view_dataset):
@@ -204,9 +218,17 @@ def build_scene_graph_pipeline(
description="Building object dataset...",
):
save_room_view(verification_output_dir, room_id, view_idx, image)
total_masks += len(masks)
total_raw_masks += len(masks)
for mask_idx, mask in enumerate(masks):
ranked_masks = rank_masks(
masks,
image_shape=(image.height, image.width),
config=mask_scoring_config,
max_masks=sam_max_masks,
)
total_masks += len(ranked_masks)
for mask_idx, mask in enumerate(ranked_masks):
masked_image = extract_masked_region(image, mask["segment"])
object_dataset.append(
(room_id, view_idx, mask_idx, mask["bbox"], masked_image)
@@ -268,7 +290,8 @@ def build_scene_graph_pipeline(
)
print(f"Total objects created: {len(scene_graph.objects)}")
print(f"Total processed masks: {total_masks}")
print(f"Total raw masks from SAM: {total_raw_masks}")
print(f"Total kept masks after ranking: {total_masks}")
print(f"Saved object images to: {verification_output_dir}")
return (scene_graph,)

4
pyproject.toml
View File

@@ -41,6 +41,10 @@ dev = [
python = "/workspace/envs/mini-nav"
root = ["./mini-nav", "./notebooks"]
[tool.basedpyright]
include = ["mini-nav", "notebooks"]
exclude = ["**/node_modules", "**/__pycache__"]
[tool.marimo.runtime]
pythonpath = ["mini-nav"]