feat(notebooks): add proposal segmentation notebook and enhance verification

notebooks/proposal_segament.py

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

notebooks/verification.py

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

pyproject.toml

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