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, hamming_distance
    from configs import cfg_manager
    from scenegraph import ObjectNode, RoomNode, SimpleSceneGraph
    from simulator import (
        HabitatSimulatorConfig,
        TopDownSceneElements,
        collect_room_views_by_room,
        create_habitat_simulator,
        render_topdown_scene_map,
    )
    from utils.image import numpy_to_pil

    return (
        HashPipeline,
        HabitatSimulatorConfig,
        ObjectNode,
        RoomNode,
        SimpleSceneGraph,
        TopDownSceneElements,
        cfg_manager,
        collect_room_views_by_room,
        create_habitat_simulator,
        hamming_distance,
        numpy_to_pil,
        render_topdown_scene_map,
    )


@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 = 512
    _num_rooms = 4
    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),
            )
        )

    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,
    render_topdown_scene_map,
    room_nodes,
    sim,
):
    from habitat.utils.visualizations import maps
    from matplotlib import pyplot as plt

    render_topdown_scene_map(
        pathfinder=sim.pathfinder,
        elements=TopDownSceneElements(room_nodes=room_nodes),
        meters_per_pixel=meters_per_pixel,
        maps_module=maps,
        plt_module=plt,
    )
    return


@app.cell
def pipeline_init(HashPipeline):
    pipeline = HashPipeline(
        dino_model="facebook/dinov2-large",
        sam_model="facebook/sam2.1-hiera-large",
        hash_bits=512,
    )
    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(
    Image,
    ObjectNode,
    SimpleSceneGraph,
    cfg_manager,
    np,
    pipeline,
    room_nodes,
    room_view_dataset,
    torch,
):
    """Build scene graph using step-by-step pipeline to capture cropped images."""
    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]

    # Step 1: Detect objects.
    _text_labels = ["object"]
    _detections = pipeline.detect_batch(_images, _text_labels)

    # Step 2: Segment with SAM.
    _bboxes_per_image = [[_d["bbox"] for _d in _dets] for _dets in _detections]
    _masks = pipeline.segment_batch(_images, _bboxes_per_image)

    # Step 3: Filter masks.
    _filtered = pipeline.filter_batch(_images, _masks)

    # Step 4: Crop images.
    _cropped_images = pipeline.crop_batch(_filtered, _masks, _detections)

    # Step 5: Extract DINO features and compress to hash.
    _batch_size = 32
    _all_bits = []
    for _i in range(0, len(_cropped_images), _batch_size):
        _batch = _cropped_images[_i : _i + _batch_size]
        _tokens = pipeline.extract_dino_batch(_batch)
        _bits = pipeline.compress_batch(_tokens)
        _all_bits.append(_bits)

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

    # Step 6: Create ObjectNodes and save cropped images.
    for _idx, (_cropped, _hash_bits) in enumerate(zip(_cropped_images, hash_tensor)):
        _room_id, _view_idx = _metadata[_idx]
        _obj_id = f"obj_{_idx:04d}"

        _bits_array = _hash_bits.detach().cpu().numpy().reshape(-1)
        _bits_binary = (_bits_array > 0).astype(np.uint8)
        _hash_bytes = np.packbits(_bits_binary).tobytes()

        object_images[_obj_id] = _cropped
        _cropped.save(output_dir / f"{_obj_id}.png")

        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,
        )

    print(f"Created {len(scene_graph.objects)} objects")
    print(f"Saved cropped images to: {output_dir}")

    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,
    hamming_distance,
    np,
    object_images,
    pipeline,
    scene_graph,
    torch,
):
    import io

    query_result = None
    query_cropped = None
    top_matches = []

    if file_upload.value:
        _query_image = Image.open(io.BytesIO(file_upload.contents())).convert("RGB")

        # Step-by-step processing to get cropped query image.
        _text_labels = ["object"]
        _detections = pipeline.detect_batch([_query_image], _text_labels)
        _bboxes = [[_d["bbox"] for _d in _dets] for _dets in _detections]
        _masks = pipeline.segment_batch([_query_image], _bboxes)
        _filtered = pipeline.filter_batch([_query_image], _masks)
        _cropped = pipeline.crop_batch(_filtered, _masks, _detections)

        _tokens = pipeline.extract_dino_batch(_cropped)
        _query_bits = pipeline.compress_batch(_tokens)

        if _query_bits.numel() > 0:
            query_cropped = _cropped[0]
            _query_tensor = _query_bits[0].int()

            _obj_ids = list(scene_graph.objects.keys())
            _obj_hashes = []
            for _obj_id in _obj_ids:
                _obj = scene_graph.objects[_obj_id]
                _bits = np.unpackbits(np.frombuffer(_obj.visual_hash, dtype=np.uint8))[
                    : pipeline.hash_bits
                ]
                _obj_hashes.append(_bits)

            if _obj_hashes:
                _db_tensor = torch.tensor(np.array(_obj_hashes), dtype=torch.int32)
                _db_tensor = _db_tensor.to(_query_tensor.device)

                _distances = hamming_distance(_query_tensor.unsqueeze(0), _db_tensor)
                _distances = _distances.squeeze(0).cpu().numpy()

                _top_k = min(5, len(_obj_ids))
                _top_indices = np.argsort(_distances)[:_top_k]

                top_matches = [
                    {
                        "obj_id": _obj_ids[_i],
                        "distance": int(_distances[_i]),
                        "similarity": 1.0 - _distances[_i] / float(pipeline.hash_bits),
                    }
                    for _i in _top_indices
                ]

                query_result = {
                    "query_cropped": query_cropped,
                    "top_matches": top_matches,
                }

    return query_cropped, query_result, top_matches


@app.cell
def display_results(mo, object_images, query_cropped, query_result, top_matches):
    if query_result is None:
        mo.md("No query results yet. Upload an image above.")
    else:
        _result_items = []

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

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

            if _obj_img:
                _result_items.append(
                    mo.vstack(
                        [
                            mo.md(f"**{_obj_id}**"),
                            mo.image(_obj_img),
                            mo.md(f"Distance: {_dist}"),
                            mo.md(f"Similarity: {_sim:.2%}"),
                        ],
                        align="center",
                    )
                )

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

    return


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