diff --git a/notebooks/proposal_segament.py b/notebooks/proposal_segament.py new file mode 100644 index 0000000..013dd32 --- /dev/null +++ b/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() diff --git a/notebooks/verification.py b/notebooks/verification.py index 402cfcf..ed7e70f 100644 --- a/notebooks/verification.py +++ b/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,) diff --git a/pyproject.toml b/pyproject.toml index 4ce0d3d..a4905a6 100644 --- a/pyproject.toml +++ b/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"]