refactor(compressors): migrate to centralized model loaders

- Refactor model_loader.py: improve return type annotations from tuple[Any, Any] to tuple[AutoImageProcessor, AutoModel]
- Refactor proposal/core.py: add input validation for mask array dimensionality, handle 2D masks and batch dimensions gracefully
- Refactor proposal_segament.ipynb: replace inline model loading with centralized load_owlv2_model() and load_sam_model() functions, use batched detect_objects_batch() and generate_proposals_batch() APIs

mini-nav/compressors/model_loader.py

@@ -34,7 +34,7 @@ def load_sam_model(
 
 def load_dino_model(
     model_name: str = "facebook/dinov2-large",
-) -> tuple[Any, Any]:
+) -> tuple[AutoImageProcessor, AutoModel]:
     device = get_device()
 
     processor = AutoImageProcessor.from_pretrained(model_name)

mini-nav/compressors/proposal/core.py

@@ -205,16 +205,14 @@ def _masks_to_proposals(masks: Any) -> list[dict[str, Any]]:
     if mask_array is None:
         return []
 
-    # Ensure 3D: [num_masks, H, W]
+    if mask_array.ndim < 2:
-    if mask_array.ndim == 2:
-        mask_array = np.expand_dims(mask_array, axis=0)
-
-    if mask_array.ndim != 3:
         return []
 
-    # Remove batch dim if present: [1, num_masks, H, W] → [num_masks, H, W]
+    if mask_array.ndim == 2:
-    if mask_array.ndim == 3 and mask_array.shape[0] == 1:
+        mask_array = np.expand_dims(mask_array, axis=0)
-        mask_array = mask_array[0]
+    else:
+        height, width = mask_array.shape[-2], mask_array.shape[-1]
+        mask_array = mask_array.reshape(-1, height, width)
 
     proposals: list[dict[str, Any]] = []
     for single_mask in mask_array:

notebooks/proposal_segament.py

@@ -8,13 +8,9 @@ app = marimo.App()
 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
+    return Image, ImageDraw, ImageFont, mo, np
 
 
 @app.cell
@@ -67,15 +63,6 @@ def _(np):
 
 @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(
@@ -88,30 +75,17 @@ def _(agent, room_nodes, sim, views_per_room):
 
 
 @app.cell
-def _(ImageDraw, ImageFont, all_room_views, mo, torch):
+def _(ImageDraw, ImageFont, all_room_views, mo):
-    from transformers import Owlv2Processor, Owlv2ForObjectDetection
+    from compressors.model_loader import load_owlv2_model
-    from utils import get_device, numpy_to_pil
+    from compressors.proposal.core import detect_objects_batch
+    from utils.common import get_device
+    from utils.image import numpy_to_pil
 
     device = get_device()
 
-    # 你真正想控制的显示阈值
     score_threshold = 0.25
-
+    postprocess_threshold = 0.1
-    # 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",
@@ -119,40 +93,26 @@ def _(ImageDraw, ImageFont, all_room_views, mo, torch):
         "a table",
         "a vase",
         "a painting",
-            "a window"
+        "a window",
-        ]
     ]
 
-    # 4. 推理
+    owl_processor, owl_model = load_owlv2_model(
-    inputs = processor(text=text_labels, images=image, return_tensors="pt").to(
+        model_name="google/owlv2-base-patch16-ensemble"
-        device
     )
 
-    with torch.no_grad():
+    image = numpy_to_pil(all_room_views["room_00"][1])
-        outputs = model(**inputs)
 
-    # 5. 后处理
+    detection_batch = detect_objects_batch(
-    target_sizes = torch.tensor([(image.height, image.width)])
+        model=owl_model,
-    results = processor.post_process_grounded_object_detection(
+        processor=owl_processor,
-        outputs=outputs,
+        images=[image],
-        target_sizes=target_sizes,
+        text_labels_per_image=[text_labels],
-        threshold=0.1,  # 模型后处理的第一层粗筛
+        score_threshold=score_threshold,
-        text_labels=text_labels,
+        postprocess_threshold=postprocess_threshold,
     )
+    detections = detection_batch[0] if detection_batch else []
+    filtered_items = [(det["bbox"], det["score"], det["label"]) for det in detections]
 
-    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)
 
@@ -162,13 +122,11 @@ def _(ImageDraw, ImageFont, all_room_views, mo, torch):
         font = ImageFont.load_default()
 
     for _box, _score, _text_label in filtered_items:
-        _x1, _y1, _x2, _y2 = [float(v) for v in _box.tolist()]
+        _x1, _y1, _x2, _y2 = [float(v) for v in _box]
-        _label = f"{_text_label}: {_score.item():.3f}"
+        _label = f"{_text_label}: {_score:.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:
@@ -183,9 +141,9 @@ def _(ImageDraw, ImageFont, all_room_views, mo, torch):
     # 8. 结果文本
     detection_lines = []
     for _box, _score, _text_label in filtered_items:
-        box_rounded = [round(_v, 2) for _v in _box.tolist()]
+        box_rounded = [round(_v, 2) for _v in _box]
         detection_lines.append(
-            f"- {_text_label}: score={_score.item():.3f}, box={box_rounded}"
+            f"- {_text_label}: score={_score:.3f}, box={box_rounded}"
         )
 
     if not detection_lines:
@@ -193,11 +151,12 @@ def _(ImageDraw, ImageFont, all_room_views, mo, torch):
     else:
         detection_text = "\n".join(detection_lines)
 
-    # 9. marimo 输出
     mo.vstack(
         [
             mo.md(
-                f"## OWLv2 检测可视化结果\n\n过滤阈值：`score >= {score_threshold:.2f}`"
+                "## OWLv2 检测可视化结果"
+                f"\n\ndevice: `{device}`"
+                f"\n\n过滤阈值：`score >= {score_threshold:.2f}`"
             ),
             mo.image(_vis_image, width=700),
             mo.md(detection_text),
@@ -207,90 +166,27 @@ def _(ImageDraw, ImageFont, all_room_views, mo, torch):
 
 
 @app.cell
-def _(Image, ImageDraw, device, filtered_items, image, mo, np, torch):
+def _(Image, ImageDraw, device, filtered_items, image, mo, np):
-    from transformers import Sam2Processor, Sam2Model
+    from compressors.model_loader import load_sam_model
+    from compressors.proposal.core import generate_proposals_batch
 
-    # -----------------------------
+    sam2_processor, sam2_model = load_sam_model(
-    # 0. 加载 SAM2
+        model_name="facebook/sam2.1-hiera-large"
-    # -----------------------------
-    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 个
+    input_boxes = [[box for box, _score, _text_label in filtered_items]]
-        # 形状通常是 [num_objects, num_masks, H, W]
+    proposal_batch = generate_proposals_batch(
-        # 由于 multimask_output=False，num_masks 通常为 1
+        model=sam2_model,
-        masks_for_image = all_masks[0]
+        processor=sam2_processor,
+        images=[image],
+        bboxes_per_image=input_boxes,
+    )
+    proposals = proposal_batch[0] if proposal_batch else []
 
-        # 兼容处理：压掉单 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))
+    _vis_image = base_rgba.copy()
-        overlay_draw = ImageDraw.Draw(overlay)
+    summary_lines = []
 
-        # 准备一个可重复的颜色列表
     colors = [
         (255, 0, 0, 90),
         (0, 255, 0, 90),
@@ -302,39 +198,23 @@ def _(Image, ImageDraw, device, filtered_items, image, mo, np, torch):
         (128, 0, 255, 90),
     ]
 
-        # 先画 mask，再画 box
+    for _idx, ((_box, _score, _text_label), proposal) in enumerate(
-        _vis_image = base_rgba.copy()
+        zip(filtered_items, proposals)
-
-        for _idx, ((_box, _score, _text_label), mask) in enumerate(
-            zip(filtered_items, masks_for_image)
     ):
+        mask_np = proposal["segment"]
         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 (_box, _score, _text_label), proposal in zip(filtered_items, proposals):
-        for _idx, ((_box, _score, _text_label), mask) in enumerate(
+        _x1, _y1, _x2, _y2 = [float(v) for v in _box]
-            zip(filtered_items, masks_for_image)
+        mask_area = int(proposal["area"])
-        ):
+        _label = f"{_text_label} | owl={_score:.3f} | mask_area={mask_area}"
-            _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)
 
@@ -348,32 +228,24 @@ def _(Image, ImageDraw, device, filtered_items, image, mo, np, torch):
             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}"
+            f"- {_text_label}: owl_score={_score:.3f}, mask_area={mask_area}"
         )
 
+    if not filtered_items:
+        summary_text = (
+            "没有可用于分割的检测框，请先降低 OWLv2 的 score_threshold 或检查检测结果。"
+        )
+    elif not summary_lines:
+        summary_text = "没有生成任何 mask"
+    else:
+        summary_text = "\n".join(summary_lines)
 
     mo.vstack(
         [
-            mo.md("## SAM2 分割可视化结果"),
+            mo.md(f"## SAM2 分割可视化结果\n\ndevice: `{device}`"),
             mo.image(_vis_image, width=700),
-            mo.md(
+            mo.md(summary_text),
-                "\n".join(summary_lines) if summary_lines else "没有生成任何 mask"
-            ),
         ]
     )
     return