# /// 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( ObjectNode, SimpleSceneGraph, cfg_manager, np, pipeline, room_nodes, room_view_dataset, ): 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] _text_labels = ["object"] _output = pipeline.process_batch(_images, _text_labels, batch_size=32) _cropped_images = _output.cropped_images hash_tensor = _output.hash_bits # 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, ) _fallback_count = sum( 1 for _meta in _output.debug_meta if _meta["fallback_reason"] is not None ) print(f"Created {len(scene_graph.objects)} objects") print(f"Saved cropped images to: {output_dir}") print(f"Fallback frames: {_fallback_count}/{len(_output.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, 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") _text_labels = ["object"] _output = pipeline.process_batch([_query_image], _text_labels, batch_size=1) _query_bits = _output.hash_bits if _query_bits.numel() > 0: query_cropped = _output.cropped_images[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()