feat(compressors): refactor pipeline with FramePacket dataclass and unified process_batch

- Add FramePacket dataclass to encapsulate per-image pipeline state
- Rename internal methods with underscore prefix convention
- Replace separate filter_batch/crop_batch with unified process_batch method
- Update notebook to use new HashPipeline API
This commit is contained in:
2026-04-04 19:55:36 +08:00
parent 94ed05a039
commit 3638ffdb8d
3 changed files with 450 additions and 243 deletions

View File

@@ -76,8 +76,7 @@ def _(agent, room_nodes, sim, views_per_room):
@app.cell
def _(ImageDraw, ImageFont, all_room_views, mo):
from compressors.model_loader import load_owlv2_model
from compressors.proposal.core import detect_objects_batch
from compressors import HashPipeline
from utils.common import get_device
from utils.image import numpy_to_pil
@@ -96,22 +95,25 @@ def _(ImageDraw, ImageFont, all_room_views, mo):
"a window",
]
owl_processor, owl_model = load_owlv2_model(
model_name="google/owlv2-base-patch16-ensemble"
pipeline = HashPipeline(
owlv2_model="google/owlv2-base-patch16-ensemble",
score_threshold=score_threshold,
postprocess_threshold=postprocess_threshold,
)
image = numpy_to_pil(all_room_views["room_00"][1])
detection_batch = detect_objects_batch(
model=owl_model,
processor=owl_processor,
output = pipeline.process_batch(
images=[image],
text_labels_per_image=[text_labels],
score_threshold=score_threshold,
postprocess_threshold=postprocess_threshold,
text_labels=text_labels,
batch_size=1,
return_debug_details=True,
)
detections = detection_batch[0] if detection_batch else []
filtered_items = [(det["bbox"], det["score"], det["label"]) for det in detections]
meta = output.debug_meta[0] if output.debug_meta else {}
boxes = meta.get("boxes_xyxy", [])
scores = meta.get("scores", [])
labels = meta.get("labels", [])
filtered_items = list(zip(boxes, scores, labels, strict=False))
_vis_image = image.copy()
_draw = ImageDraw.Draw(_vis_image)
@@ -162,32 +164,56 @@ def _(ImageDraw, ImageFont, all_room_views, mo):
mo.md(detection_text),
]
)
return device, filtered_items, image
return device, filtered_items, image, meta
@app.cell
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
def _(meta):
proposals = meta.get("masks", [])
return (proposals,)
sam2_processor, sam2_model = load_sam_model(
model_name="facebook/sam2.1-hiera-large"
@app.cell
def _(Image, ImageDraw, filtered_items, image, mo, np, proposals):
from compressors.filter import (
MaskScoringConfig,
compute_mask_features,
score_mask,
should_reject_mask,
)
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 []
image_shape = (image.height, image.width)
config = MaskScoringConfig()
base_rgba = image.convert("RGBA")
_vis_image = base_rgba.copy()
summary_lines = []
_kept = []
_rejected = []
for _idx, proposal in enumerate(proposals):
_feat = compute_mask_features(proposal, image_shape)
_is_rejected = should_reject_mask(_feat, config)
_score = score_mask(proposal, image_shape, config)
_owl_label = (
filtered_items[_idx][2] if _idx < len(filtered_items) else f"obj_{_idx}"
)
_owl_score = filtered_items[_idx][1] if _idx < len(filtered_items) else 0.0
_owl_bbox = (
filtered_items[_idx][0] if _idx < len(filtered_items) else [0, 0, 0, 0]
)
colors = [
_entry = {
"idx": _idx,
"proposal": proposal,
"features": _feat,
"mask_score": _score,
"owl_label": _owl_label,
"owl_score": _owl_score,
"owl_bbox": _owl_bbox,
}
if _is_rejected:
_rejected.append(_entry)
else:
_kept.append(_entry)
_colors = [
(255, 0, 0, 90),
(0, 255, 0, 90),
(0, 0, 255, 90),
@@ -198,54 +224,143 @@ def _(Image, ImageDraw, device, filtered_items, image, mo, np):
(128, 0, 255, 90),
]
for _idx, ((_box, _score, _text_label), proposal) in enumerate(
zip(filtered_items, proposals)
):
mask_np = proposal["segment"]
color = colors[_idx % len(colors)]
def _overlay_masks(base_img, entries, border_color, show_score=False):
"""Overlay masks on image and draw bounding boxes with labels."""
_rgba = base_img.convert("RGBA").copy()
for _e in entries:
_p = _e["proposal"]
_c = _colors[_e["idx"] % len(_colors)]
_mask_rgba = np.zeros((base_img.height, base_img.width, 4), dtype=np.uint8)
_mask_rgba[_p["segment"]] = _c
_rgba = Image.alpha_composite(
_rgba, Image.fromarray(_mask_rgba, mode="RGBA")
)
mask_rgba = np.zeros((image.height, image.width, 4), dtype=np.uint8)
mask_rgba[mask_np] = color
_draw = ImageDraw.Draw(_rgba)
for _e in entries:
_x1, _y1, _x2, _y2 = [float(v) for v in _e["owl_bbox"]]
_mask_area = int(_e["proposal"]["area"])
if show_score:
_label = f"{_e['owl_label']} | score={_e['mask_score']:.2f} | area={_mask_area}"
else:
_label = f"{_e['owl_label']} | area={_mask_area}"
_draw.rectangle([_x1, _y1, _x2, _y2], outline=border_color, width=3)
mask_img = Image.fromarray(mask_rgba, mode="RGBA")
_vis_image = Image.alpha_composite(_vis_image, mask_img)
try:
_tb = _draw.textbbox((_x1, _y1), _label)
except Exception:
_tb = (_x1, _y1, _x1 + 200, _y1 + 20)
_draw.rectangle(
[_tb[0], max(0, _tb[1] - 2), _tb[2] + 4, _tb[3] + 2],
fill=border_color,
)
_draw.text((_x1 + 2, max(0, _y1)), _label, fill="white")
return _rgba
_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}"
_all_entries = _kept + _rejected
_all_entries.sort(key=lambda e: e["idx"])
_draw.rectangle([_x1, _y1, _x2, _y2], outline=(255, 0, 0, 255), width=3)
_before_img = _overlay_masks(image, _all_entries, "red", show_score=False)
_after_img = _overlay_masks(image, _kept, (0, 180, 0), show_score=True)
try:
_tx1, _ty1, _tx2, _ty2 = _draw.textbbox((_x1, _y1), _label)
except Exception:
_tx1, _ty1, _tx2, _ty2 = _x1, _y1, _x1 + 220, _y1 + 20
_draw_after = ImageDraw.Draw(_after_img)
for _e in _rejected:
_x1, _y1, _x2, _y2 = [float(v) for v in _e["owl_bbox"]]
_feat = _e["features"]
_reason_parts = []
if _feat.area_ratio < config.min_area_ratio:
_reason_parts.append(f"area_ratio={_feat.area_ratio:.4f}<min")
if _feat.area_ratio > config.max_area_ratio:
_reason_parts.append(f"area_ratio={_feat.area_ratio:.4f}>max")
_aspect = max(_feat.aspect_ratio, 1.0 / max(_feat.aspect_ratio, 1e-6))
if _aspect > config.max_aspect_ratio:
_reason_parts.append(f"aspect={_aspect:.1f}>max")
if _feat.fill_ratio < config.min_fill_ratio_hard:
_reason_parts.append(f"fill={_feat.fill_ratio:.3f}<min")
if (
_feat.num_components > config.max_components
and _feat.largest_component_ratio < config.min_largest_component_ratio
):
_reason_parts.append(f"fragments={_feat.num_components}")
if _feat.touch_edge_count >= config.reject_edge_touch_count:
_reason_parts.append(f"edge_touch={_feat.touch_edge_count}")
if (
_feat.touch_edge_count >= config.reject_large_edge_touch_count
and _feat.area_ratio > config.reject_large_edge_area_ratio
):
_reason_parts.append("edge+large")
_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}"
_reason = ", ".join(_reason_parts) if _reason_parts else "unknown"
_label = f"X {_e['owl_label']}: {_reason}"
_dash_len = 8
_gap_len = 4
for _seg_start in range(int(_x1), int(_x2), _dash_len + _gap_len):
_seg_end = min(_seg_start + _dash_len, int(_x2))
_draw_after.line([(_seg_start, _y1), (_seg_end, _y1)], fill="red", width=2)
_draw_after.line([(_seg_start, _y2), (_seg_end, _y2)], fill="red", width=2)
for _seg_start in range(int(_y1), int(_y2), _dash_len + _gap_len):
_seg_end = min(_seg_start + _dash_len, int(_y2))
_draw_after.line([(_x1, _seg_start), (_x1, _seg_end)], fill="red", width=2)
_draw_after.line([(_x2, _seg_start), (_x2, _seg_end)], fill="red", width=2)
_draw_after.text((_x1 + 2, max(0, _y1 - 18)), _label, fill="red")
_total = len(proposals)
_kept_count = len(_kept)
_rej_count = len(_rejected)
_rej_detail_lines = []
for _e in _rejected:
_f = _e["features"]
_rej_detail_lines.append(
f" - **{_e['owl_label']}** (idx={_e['idx']}): "
f"area_ratio={_f.area_ratio:.4f}, fill_ratio={_f.fill_ratio:.3f}, "
f"aspect_ratio={_f.aspect_ratio:.1f}, touch_edge={_f.touch_edge_count}, "
f"components={_f.num_components}, mask_score={_e['mask_score']:.3f}"
)
if not filtered_items:
summary_text = (
"没有可用于分割的检测框,请先降低 OWLv2 的 score_threshold 或检查检测结果。"
_kept_detail_lines = []
for _e in _kept:
_f = _e["features"]
_kept_detail_lines.append(
f" - **{_e['owl_label']}** (idx={_e['idx']}): "
f"area_ratio={_f.area_ratio:.4f}, mask_score={_e['mask_score']:.3f}"
)
elif not summary_lines:
summary_text = "没有生成任何 mask"
else:
summary_text = "\n".join(summary_lines)
_detail_parts = []
if _kept_detail_lines:
_detail_parts.append("**保留的 mask**\n" + "\n".join(_kept_detail_lines))
if _rej_detail_lines:
_detail_parts.append("**过滤掉的 mask**\n" + "\n".join(_rej_detail_lines))
_detail_text = "\n\n".join(_detail_parts) if _detail_parts else "无 mask 数据"
mo.vstack(
[
mo.md(f"## SAM2 分割可视化结果\n\ndevice: `{device}`"),
mo.image(_vis_image, width=700),
mo.md(summary_text),
mo.md(
"## Mask 过滤对比"
f"\n\n{_total} 个 mask → 保留 **{_kept_count}** 个,过滤掉 **{_rej_count}** 个"
),
mo.hstack(
[
mo.vstack(
[
mo.md(f"### 过滤前({_total} 个)"),
mo.image(_before_img, width=480),
]
),
mo.vstack(
[
mo.md(
f"### 过滤后({_kept_count} 个保留,{_rej_count} 个过滤)"
),
mo.image(_after_img, width=480),
]
),
]
),
mo.md(_detail_text),
]
)
return