refactor(visualization): enhance proposal filtering visualization

notebooks/proposal_segament.py

@@ -22,7 +22,6 @@ def _(np):
         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
@@ -52,13 +51,12 @@ def _(np):
             )
         )
 
-    render_topdown_scene_map(
+    topdown_image = 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
+    return agent, room_nodes, sim, topdown_image, views_per_room
 
 
 @app.cell
@@ -75,7 +73,7 @@ def _(agent, room_nodes, sim, views_per_room):
 
 
 @app.cell
-def _(ImageDraw, ImageFont, all_room_views, mo):
+def _(ImageDraw, ImageFont, all_room_views, mo, topdown_image):
     from compressors import HashPipeline
     from utils.common import get_device
     from utils.image import numpy_to_pil
@@ -113,7 +111,40 @@ def _(ImageDraw, ImageFont, all_room_views, mo):
     boxes = meta.get("boxes_xyxy", [])
     scores = meta.get("scores", [])
     labels = meta.get("labels", [])
-    filtered_items = list(zip(boxes, scores, labels, strict=False))
+    masks = meta.get("masks", [])
+    pipeline_selected_indices = meta.get("selected_indices", [])
+    dropped_indices = meta.get("dropped_indices", [])
+
+    proposal_items = []
+    for idx, proposal_data in enumerate(masks):
+        detection_box = boxes[idx] if idx < len(boxes) else [0.0, 0.0, 0.0, 0.0]
+        detection_score = scores[idx] if idx < len(scores) else 0.0
+        detection_label = labels[idx] if idx < len(labels) else f"obj_{idx}"
+        proposal_items.append(
+            {
+                "idx": idx,
+                "proposal": proposal_data,
+                "bbox": detection_box,
+                "score": detection_score,
+                "label": detection_label,
+                "is_selected": idx in pipeline_selected_indices,
+                "is_dropped": idx in dropped_indices,
+            }
+        )
+
+    detected_items = [
+        {
+            "idx": idx,
+            "bbox": box,
+            "score": score,
+            "label": label,
+            "has_mask": idx < len(masks),
+            "is_selected": idx in pipeline_selected_indices,
+        }
+        for idx, (box, score, label) in enumerate(
+            zip(boxes, scores, labels, strict=False)
+        )
+    ]
 
     _vis_image = image.copy()
     _draw = ImageDraw.Draw(_vis_image)
@@ -123,11 +154,12 @@ def _(ImageDraw, ImageFont, all_room_views, mo):
     except Exception:
         font = ImageFont.load_default()
 
-    for _box, _score, _text_label in filtered_items:
+    for _item in detected_items:
-        _x1, _y1, _x2, _y2 = [float(v) for v in _box]
+        _x1, _y1, _x2, _y2 = [float(v) for v in _item["bbox"]]
-        _label = f"{_text_label}: {_score:.3f}"
+        _label = f"{_item['label']}: {_item['score']:.3f}"
+        _outline_color = "green" if _item["is_selected"] else "red"
 
-        _draw.rectangle([_x1, _y1, _x2, _y2], outline="red", width=3)
+        _draw.rectangle([_x1, _y1, _x2, _y2], outline=_outline_color, width=3)
 
         try:
             _tx1, _ty1, _tx2, _ty2 = _draw.textbbox((_x1, _y1), _label, font=font)
@@ -137,15 +169,15 @@ def _(ImageDraw, ImageFont, all_room_views, mo):
             _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.rectangle(text_bg, fill=_outline_color)
         _draw.text((_x1 + 2, max(0, _y1)), _label, fill="white", font=font)
 
-    # 8. 结果文本
     detection_lines = []
-    for _box, _score, _text_label in filtered_items:
+    for _item in detected_items:
-        box_rounded = [round(_v, 2) for _v in _box]
+        box_rounded = [round(_v, 2) for _v in _item["bbox"]]
+        status = "selected" if _item["is_selected"] else "dropped"
         detection_lines.append(
-            f"- {_text_label}: score={_score:.3f}, box={box_rounded}"
+            f"- {_item['label']}: score={_item['score']:.3f}, box={box_rounded}, status={status}"
         )
 
     if not detection_lines:
@@ -159,59 +191,56 @@ def _(ImageDraw, ImageFont, all_room_views, mo):
                 "## OWLv2 检测可视化结果"
                 f"\n\ndevice: `{device}`"
                 f"\n\n过滤阈值：`score >= {score_threshold:.2f}`"
+                f"\n\n绿色框表示最终被 pipeline 保留，红色框表示未被保留"
             ),
+            mo.md("## Top-down 房间采样图"),
+            mo.image(topdown_image, width=520),
             mo.image(_vis_image, width=700),
             mo.md(detection_text),
         ]
     )
-    return device, filtered_items, image, meta
+    return detected_items, device, image, meta, proposal_items
 
 
 @app.cell
-def _(meta):
+def _(Image, ImageDraw, image, meta, mo, np, proposal_items):
-    proposals = meta.get("masks", [])
-    return (proposals,)
-
-
-@app.cell
-def _(Image, ImageDraw, filtered_items, image, mo, np, proposals):
     from compressors.filter import (
         MaskScoringConfig,
         compute_mask_features,
         score_mask,
-        should_reject_mask,
     )
 
     image_shape = (image.height, image.width)
     config = MaskScoringConfig()
+    fallback_reason = meta.get("fallback_reason")
+    selected_index_set = set(meta.get("selected_indices", []))
 
     _kept = []
     _rejected = []
-    for _idx, proposal in enumerate(proposals):
+    for _item in proposal_items:
-        _feat = compute_mask_features(proposal, image_shape)
+        _idx = _item["idx"]
-        _is_rejected = should_reject_mask(_feat, config)
+        current_proposal = _item["proposal"]
-        _score = score_mask(proposal, image_shape, config)
+        _feat = compute_mask_features(current_proposal, image_shape)
-        _owl_label = (
+        _score = score_mask(current_proposal, image_shape, config)
-            filtered_items[_idx][2] if _idx < len(filtered_items) else f"obj_{_idx}"
+        _owl_label = _item["label"]
-        )
+        _owl_score = _item["score"]
-        _owl_score = filtered_items[_idx][1] if _idx < len(filtered_items) else 0.0
+        _owl_bbox = _item["bbox"]
-        _owl_bbox = (
+        _is_selected = _idx in selected_index_set
-            filtered_items[_idx][0] if _idx < len(filtered_items) else [0, 0, 0, 0]
-        )
 
         _entry = {
             "idx": _idx,
-            "proposal": proposal,
+            "proposal": current_proposal,
             "features": _feat,
             "mask_score": _score,
             "owl_label": _owl_label,
             "owl_score": _owl_score,
             "owl_bbox": _owl_bbox,
+            "is_selected": _is_selected,
         }
-        if _is_rejected:
+        if _is_selected:
-            _rejected.append(_entry)
-        else:
             _kept.append(_entry)
+        else:
+            _rejected.append(_entry)
 
     _colors = [
         (255, 0, 0, 90),
@@ -289,6 +318,8 @@ def _(Image, ImageDraw, filtered_items, image, mo, np, proposals):
             and _feat.area_ratio > config.reject_large_edge_area_ratio
         ):
             _reason_parts.append("edge+large")
+        if _e["idx"] in meta.get("dropped_indices", []):
+            _reason_parts.append("pipeline_drop")
 
         _reason = ", ".join(_reason_parts) if _reason_parts else "unknown"
         _label = f"X {_e['owl_label']}: {_reason}"
@@ -306,7 +337,7 @@ def _(Image, ImageDraw, filtered_items, image, mo, np, proposals):
 
         _draw_after.text((_x1 + 2, max(0, _y1 - 18)), _label, fill="red")
 
-    _total = len(proposals)
+    _total = len(proposal_items)
     _kept_count = len(_kept)
     _rej_count = len(_rejected)
 
@@ -335,25 +366,27 @@ def _(Image, ImageDraw, filtered_items, image, mo, np, proposals):
         _detail_parts.append("**过滤掉的 mask：**\n" + "\n".join(_rej_detail_lines))
 
     _detail_text = "\n\n".join(_detail_parts) if _detail_parts else "无 mask 数据"
+    if fallback_reason is not None:
+        _detail_text = f"**Pipeline fallback:** `{fallback_reason}`\n\n" + _detail_text
 
     mo.vstack(
         [
             mo.md(
                 "## Mask 过滤对比"
-                f"\n\n共 {_total} 个 mask → 保留 **{_kept_count}** 个，过滤掉 **{_rej_count}** 个"
+                f"\n\n共 {_total} 个 proposal → 保留 **{_kept_count}** 个，过滤掉 **{_rej_count}** 个"
             ),
             mo.hstack(
                 [
                     mo.vstack(
                         [
-                            mo.md(f"### 过滤前（{_total} 个）"),
+                            mo.md(f"### 过滤前（全部 {_total} 个 proposal）"),
                             mo.image(_before_img, width=480),
                         ]
                     ),
                     mo.vstack(
                         [
                             mo.md(
-                                f"### 过滤后（{_kept_count} 个保留，{_rej_count} 个过滤）"
+                                f"### 过滤后（{_kept_count} 个保留，{_rej_count} 个未保留）"
                             ),
                             mo.image(_after_img, width=480),
                         ]