mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-13 04:25:32 +08:00
384 lines
12 KiB
Python
384 lines
12 KiB
Python
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()
|