diff --git a/notebooks/proposal_segament.py b/notebooks/proposal_segament.py index b5f8dcb..0bfb89a 100644 --- a/notebooks/proposal_segament.py +++ b/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: - _x1, _y1, _x2, _y2 = [float(v) for v in _box] - _label = f"{_text_label}: {_score:.3f}" + for _item in detected_items: + _x1, _y1, _x2, _y2 = [float(v) for v in _item["bbox"]] + _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: - box_rounded = [round(_v, 2) for _v in _box] + for _item in detected_items: + 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): - proposals = meta.get("masks", []) - return (proposals,) - - -@app.cell -def _(Image, ImageDraw, filtered_items, image, mo, np, proposals): +def _(Image, ImageDraw, image, meta, mo, np, proposal_items): 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): - _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] - ) + for _item in proposal_items: + _idx = _item["idx"] + current_proposal = _item["proposal"] + _feat = compute_mask_features(current_proposal, image_shape) + _score = score_mask(current_proposal, image_shape, config) + _owl_label = _item["label"] + _owl_score = _item["score"] + _owl_bbox = _item["bbox"] + _is_selected = _idx in selected_index_set _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: - _rejected.append(_entry) - else: + if _is_selected: _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), ]