Mini-Nav/notebooks/verification.py

# /// script
# requires-python = ">=3.13"
# dependencies = [
#     "marimo>=0.21.1",
#     "pyzmq>=27.1.0",
# ]
# ///

import marimo

__generated_with = "0.21.1"
app = marimo.App(width="medium", app_title="Pipeline Verification")


@app.cell
def _():
    import marimo as mo

    return (mo,)


@app.cell
def base_dependencies():
    """Basic dependencies for data processing."""
    import numpy as np
    import polars as pl
    import torch
    from PIL import Image

    return Image, np, pl, torch


@app.cell
def project_imports():
    """Project module imports using new architecture."""
    from compressors import HashPipeline
    from configs import cfg_manager
    from scenegraph import (
        ObjectNode,
        RoomNode,
        SimpleSceneGraph,
        bits_tensor_to_hash_bytes,
        query_image_against_scene_graph,
    )
    from simulator import (
        HabitatSimulatorConfig,
        TopDownSceneElements,
        collect_room_views_by_room,
        create_habitat_simulator,
        render_topdown_scene_map,
        save_object_image,
        save_room_view,
    )
    from utils.image import numpy_to_pil

    return (
        HashPipeline,
        HabitatSimulatorConfig,
        ObjectNode,
        RoomNode,
        SimpleSceneGraph,
        TopDownSceneElements,
        bits_tensor_to_hash_bytes,
        cfg_manager,
        collect_room_views_by_room,
        create_habitat_simulator,
        numpy_to_pil,
        query_image_against_scene_graph,
        render_topdown_scene_map,
        save_object_image,
        save_room_view,
    )


@app.cell
def text_labels():
    """Shared text labels for detection during graph build and query."""
    text_labels = [
        "a chair",
        "a table",
        "a sofa",
        "a cabinet",
        "a shelf",
        "a lamp",
        "a picture",
        "a window",
        "a door",
        "a plant",
    ]
    return (text_labels,)


@app.cell
def habitat_setup(HabitatSimulatorConfig, RoomNode, create_habitat_simulator, np):
    """Initialize Habitat simulator and sample room nodes."""
    _scene_path = "data/scene_datasets/habitat-test-scenes/skokloster-castle.glb"
    _image_size = 768
    _num_rooms = 5
    views_per_room = 12
    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),
            )
        )

    print("Sampled room centers:")
    for _node in room_nodes:
        print(f"  {_node.room_id}: {_node.center}")

    return agent, meters_per_pixel, room_nodes, sim, views_per_room


@app.cell
def render_topdown(
    TopDownSceneElements,
    meters_per_pixel,
    mo,
    render_topdown_scene_map,
    room_nodes,
    sim,
):
    image = render_topdown_scene_map(
        pathfinder=sim.pathfinder,
        elements=TopDownSceneElements(room_nodes=room_nodes),
        meters_per_pixel=meters_per_pixel,
    )
    mo.image(image)
    return (image,)


@app.cell
def pipeline_init(HashPipeline):
    pipeline = HashPipeline(
        dino_model="facebook/dinov2-large",
        sam_model="facebook/sam2.1-hiera-large",
        hash_bits=512,
        score_threshold=0.10,
        postprocess_threshold=0.05,
    )
    print(f"Pipeline initialized: {pipeline.hash_bits} bits")
    return (pipeline,)


@app.cell
def collect_views(
    agent,
    collect_room_views_by_room,
    numpy_to_pil,
    room_nodes,
    sim,
    views_per_room,
):
    all_room_views = collect_room_views_by_room(
        agent=agent,
        sim=sim,
        room_nodes=room_nodes,
        views_per_room=views_per_room,
    )

    # Flatten room views into (room_id, view_idx, PIL.Image) tuples.
    room_view_dataset = [
        (_room_id, _view_idx, numpy_to_pil(_rgb))
        for _room_id, _views in all_room_views.items()
        for _view_idx, _rgb in enumerate(_views)
    ]

    print(f"Collected {len(room_view_dataset)} room views")
    return all_room_views, room_view_dataset


@app.cell
def build_scene_graph(
    ObjectNode,
    SimpleSceneGraph,
    bits_tensor_to_hash_bytes,
    cfg_manager,
    mo,
    np,
    pipeline,
    room_nodes,
    room_view_dataset,
    save_object_image,
    save_room_view,
    text_labels,
    torch,
):
    scene_graph = SimpleSceneGraph(
        rooms={_room.room_id: _room for _room in room_nodes},
        objects={},
    )

    # Storage for cropped object images (for visualization).
    object_images = {}

    output_dir = cfg_manager.get().output.directory / "verification"
    output_dir.mkdir(parents=True, exist_ok=True)

    _images = [item[2] for item in room_view_dataset]
    _metadata = [(item[0], item[1]) for item in room_view_dataset]

    inference_batch_size = 4
    image_batches = [
        _images[index : index + inference_batch_size]
        for index in range(0, len(_images), inference_batch_size)
    ]

    _cropped_images = []
    debug_meta = []
    hash_batches = []
    for _batch_images in mo.status.progress_bar(
        image_batches,
        title="Running pipeline inference on room views",
        subtitle=f"Batch size {inference_batch_size} with ETA",
        completion_title="Pipeline inference finished",
        completion_subtitle=(
            f"Processed {len(_images)} room views in {len(image_batches)} batches"
        ),
        show_eta=True,
        show_rate=True,
        remove_on_exit=False,
    ):
        _batch_output = pipeline.process_batch(
            _batch_images,
            text_labels,
            batch_size=inference_batch_size,
            return_debug_details=True,
        )
        _cropped_images.extend(_batch_output.cropped_images)
        debug_meta.extend(_batch_output.debug_meta)
        if _batch_output.hash_bits.numel() > 0:
            hash_batches.append(_batch_output.hash_bits)

    if hash_batches:
        hash_tensor = torch.cat(hash_batches, dim=0)
    else:
        hash_tensor = torch.empty(
            (0, pipeline.hash_bits), dtype=torch.int32, device=pipeline.device
        )

    from collections import Counter

    _reasons = Counter(m["fallback_reason"] or "ok" for m in debug_meta)
    print(f"Fallback breakdown: {_reasons}")

    # Save original room views.
    for _room_id, _view_idx, _image in mo.status.progress_bar(
        room_view_dataset,
        title="Saving room-view snapshots",
        subtitle="Writing original room images to disk",
        completion_title="Room-view snapshots saved",
        completion_subtitle=f"Saved {len(room_view_dataset)} room views",
        show_eta=True,
        show_rate=True,
        remove_on_exit=False,
    ):
        save_room_view(output_dir, _room_id, _view_idx, _image)

    # Prefix sum: map flat crop index to (input_image_idx, mask_idx).
    _num_selected = [_m["num_selected"] for _m in debug_meta]
    assert sum(_num_selected) == len(_cropped_images), (
        f"Sum of num_selected ({sum(_num_selected)}) != cropped_images count ({len(_cropped_images)})"
    )
    _prefix_sums = [0]
    for _n in _num_selected:
        _prefix_sums.append(_prefix_sums[-1] + _n)

    _obj_counter = 0
    _total_crops = len(_cropped_images)
    object_tasks = []
    for _img_idx, _n_crops in enumerate(_num_selected):
        _room_id, _view_idx = _metadata[_img_idx]
        for _mask_idx in range(_n_crops):
            object_tasks.append((_img_idx, _room_id, _view_idx, _mask_idx, _n_crops))

    for _img_idx, _room_id, _view_idx, _mask_idx, _n_crops in mo.status.progress_bar(
        object_tasks,
        title="Building scene graph objects",
        subtitle="Preparing cropped objects and hashes with ETA",
        completion_title="Scene graph build complete",
        completion_subtitle=f"Created {_total_crops} cropped object entries",
        show_eta=True,
        show_rate=True,
        remove_on_exit=False,
    ):
        _crop_flat_idx = _prefix_sums[_img_idx] + _mask_idx
        _cropped = _cropped_images[_crop_flat_idx]
        _hash_bits = hash_tensor[_crop_flat_idx]

        _obj_id = f"{_room_id}_v{_view_idx:03d}_m{_mask_idx:02d}"

        _hash_bytes = bits_tensor_to_hash_bytes(_hash_bits)

        object_images[_obj_id] = _cropped
        save_object_image(output_dir, _room_id, _obj_id, _view_idx, _mask_idx, _cropped)

        scene_graph.objects[_obj_id] = ObjectNode(
            obj_id=_obj_id,
            room_id=_room_id,
            position=np.array([0.0, 0.0, 0.0], dtype=np.float32),
            visual_hash=_hash_bytes,
            semantic_hash=_hash_bytes,
            hit_count=1,
            last_seen_frame=_view_idx,
        )
        _obj_counter += 1

    _fallback_count = sum(
        1 for _meta in debug_meta if _meta["fallback_reason"] is not None
    )

    print(f"Created {_obj_counter} objects")
    print(f"Saved cropped images to: {output_dir}")
    print(f"Fallback frames: {_fallback_count}/{len(debug_meta)}")

    return hash_tensor, object_images, output_dir, scene_graph


@app.cell
def build_tables(pl, scene_graph):
    room_rows = [
        {
            "room_id": _room.room_id,
            "center_x": float(_room.center[0]),
            "center_y": float(_room.center[1]),
            "center_z": float(_room.center[2]),
        }
        for _room in scene_graph.rooms.values()
    ]

    object_rows = [
        {
            "obj_id": obj.obj_id,
            "room_id": obj.room_id,
            "visual_hash": obj.visual_hash.hex()[:16] + "...",
        }
        for obj in scene_graph.objects.values()
    ]

    rooms_df = pl.DataFrame(room_rows)
    objects_df = pl.DataFrame(object_rows)

    return objects_df, rooms_df


@app.cell
def upload_query(mo):
    file_upload = mo.ui.file(
        filetypes=["image/*"],
        kind="area",
        label="Upload a query image to find matching objects",
    )
    file_upload
    return (file_upload,)


@app.cell
def query_matching(
    Image,
    file_upload,
    mo,
    object_images,
    pipeline,
    query_image_against_scene_graph,
    scene_graph,
    text_labels,
):
    from io import BytesIO

    query_result = None
    query_cropped = None
    top_matches = []

    _file_contents = file_upload.contents()
    mo.stop(not _file_contents, mo.md("请先上传文件"))

    _query_image = Image.open(BytesIO(_file_contents)).convert("RGB")

    _query_results = query_image_against_scene_graph(
        image=_query_image,
        pipeline=pipeline,
        scene_graph=scene_graph,
        text_labels=text_labels,
        top_k=5,
        batch_size=1,
    )

    if _query_results:
        _best_result = max(
            _query_results,
            key=lambda result: result.matches[0].score if result.matches else -1,
        )
        query_cropped = _best_result.query_crop
        top_matches = [
            {
                "obj_id": match.obj_id,
                "distance": int(pipeline.hash_bits - match.score),
                "similarity": match.similarity,
                "hash_hex": match.hash_bytes.hex(),
            }
            for match in _best_result.matches
        ]

        query_result = {
            "query_cropped": query_cropped,
            "query_hash_hex": _best_result.query_hash.hex(),
            "top_matches": top_matches,
            "num_query_results": len(_query_results),
        }

    return query_cropped, query_result, top_matches


@app.cell
def display_results(mo, object_images, query_cropped, query_result, top_matches):
    mo.stop(not query_result, mo.md("No query results yet. Upload an image above."))

    _result_items = [
        mo.vstack(
            [
                mo.md("**Query (cropped)**"),
                mo.image(query_cropped),
            ],
            align="center",
        )
    ]

    for _match in top_matches:
        _obj_id = _match["obj_id"]
        _obj_img = object_images.get(_obj_id)

        if _obj_img is not None:
            _result_items.append(
                mo.vstack(
                    [
                        mo.md(f"**{_obj_id}**"),
                        mo.image(_obj_img),
                        mo.md(f"Distance: {_match['distance']}"),
                        mo.md(f"Similarity: {_match['similarity']:.2%}"),
                    ],
                    align="center",
                )
            )

    mo.vstack(_result_items, justify="center", gap=2)

    return


if __name__ == "__main__":
    app.run()