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

refactor(compressors): migrate to centralized model loaders

Browse Source 
- 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
This commit is contained in:
SikongJueluo
2026-04-02 21:23:06 +08:00
parent af0531a5eb
commit 4918b654e7
3 changed files with 111 additions and 241 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
mini-nav/compressors/model_loader.py
View File

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

14
mini-nav/compressors/proposal/core.py
View File

@@ -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:

336
notebooks/proposal_segament.py
View File

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