From 4918b654e79ecc576d4b5e862773909e27ba944c Mon Sep 17 00:00:00 2001
From: SikongJueluo <selfconfusion@gmail.com>
Date: Thu, 2 Apr 2026 21:23:06 +0800
Subject: [PATCH] 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  |   2 +-
 mini-nav/compressors/proposal/core.py |  14 +-
 notebooks/proposal_segament.py        | 336 ++++++++------------------
 3 files changed, 111 insertions(+), 241 deletions(-)

diff --git a/mini-nav/compressors/model_loader.py b/mini-nav/compressors/model_loader.py
index 5365031..e8b4b59 100644
--- a/mini-nav/compressors/model_loader.py
+++ b/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)
diff --git a/mini-nav/compressors/proposal/core.py b/mini-nav/compressors/proposal/core.py
index d1c1e2f..00bad8b 100644
--- a/mini-nav/compressors/proposal/core.py
+++ b/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:
-        mask_array = np.expand_dims(mask_array, axis=0)
-
-    if mask_array.ndim != 3:
+    if mask_array.ndim < 2:
         return []
 
-    # Remove batch dim if present: [1, num_masks, H, W] → [num_masks, H, W]
-    if mask_array.ndim == 3 and mask_array.shape[0] == 1:
-        mask_array = mask_array[0]
+    if mask_array.ndim == 2:
+        mask_array = np.expand_dims(mask_array, axis=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:
diff --git a/notebooks/proposal_segament.py b/notebooks/proposal_segament.py
index 013dd32..188c594 100644
--- a/notebooks/proposal_segament.py
+++ b/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,71 +75,44 @@ def _(agent, room_nodes, sim, views_per_room):
 
 
 @app.cell
-def _(ImageDraw, ImageFont, all_room_views, mo, torch):
-    from transformers import Owlv2Processor, Owlv2ForObjectDetection
-    from utils import get_device, numpy_to_pil
+def _(ImageDraw, ImageFont, all_room_views, mo):
+    from compressors.model_loader import load_owlv2_model
+    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
+    text_labels = [
+        "a cat",
+        "a dog",
+        "a TV remote control",
+        "a chair",
+        "a table",
+        "a vase",
+        "a painting",
+        "a window",
+    ]
 
-    # 1. 加载模型
-    processor = Owlv2Processor.from_pretrained(
-        "google/owlv2-base-patch16-ensemble"
+    owl_processor, owl_model = load_owlv2_model(
+        model_name="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
+    detection_batch = detect_objects_batch(
+        model=owl_model,
+        processor=owl_processor,
+        images=[image],
+        text_labels_per_image=[text_labels],
+        score_threshold=score_threshold,
+        postprocess_threshold=postprocess_threshold,
     )
+    detections = detection_batch[0] if detection_batch else []
+    filtered_items = [(det["bbox"], det["score"], det["label"]) for det in detections]
 
-    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)
 
@@ -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()]
-        _label = f"{_text_label}: {_score.item():.3f}"
+        _x1, _y1, _x2, _y2 = [float(v) for v in _box]
+        _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,173 +166,86 @@ def _(ImageDraw, ImageFont, all_room_views, mo, torch):
 
 
 @app.cell
-def _(Image, ImageDraw, device, filtered_items, image, mo, np, torch):
-    from transformers import Sam2Processor, Sam2Model
+def _(Image, ImageDraw, device, filtered_items, image, mo, np):
+    from compressors.model_loader import load_sam_model
+    from compressors.proposal.core import generate_proposals_batch
 
-    # -----------------------------
-    # 0. 加载 SAM2
-    # -----------------------------
-    sam2_model_id = "facebook/sam2.1-hiera-large"
+    sam2_processor, sam2_model = load_sam_model(
+        model_name="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()
+    input_boxes = [[box for box, _score, _text_label in filtered_items]]
+    proposal_batch = generate_proposals_batch(
+        model=sam2_model,
+        processor=sam2_processor,
+        images=[image],
+        bboxes_per_image=input_boxes,
+    )
+    proposals = proposal_batch[0] if proposal_batch else []
 
-    # -----------------------------
-    # 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
-        ]
+    base_rgba = image.convert("RGBA")
+    _vis_image = base_rgba.copy()
+    summary_lines = []
+
+    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),
     ]
 
-    # 没有框就直接提前返回
-    if len(input_boxes[0]) == 0:
-        mo.vstack(
-            [
-                mo.md("## SAM2 分割结果"),
-                mo.md(
-                    "没有可用于分割的检测框，请先降低 OWLv2 的 score_threshold 或检查检测结果。"
-                ),
-            ]
+    for _idx, ((_box, _score, _text_label), proposal) in enumerate(
+        zip(filtered_items, proposals)
+    ):
+        mask_np = proposal["segment"]
+        color = colors[_idx % len(colors)]
+
+        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):
+        _x1, _y1, _x2, _y2 = [float(v) for v in _box]
+        mask_area = int(proposal["area"])
+        _label = f"{_text_label} | owl={_score:.3f} | mask_area={mask_area}"
+
+        _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")
+        summary_lines.append(
+            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:
-        # -----------------------------
-        # 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}"
-            )
-
+        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(
-                "\n".join(summary_lines) if summary_lines else "没有生成任何 mask"
-            ),
+            mo.md(summary_text),
         ]
     )
     return