mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-12 20:15:31 +08:00
SikongJueluo
ba96cec406
- Export ImageSceneGraphQueryResult and query_image_against_scene_graph from scenegraph module - Replace inline hamming-distance-based image matching with dedicated query_image_against_scene_graph function - Improve top_matches structure by extracting similarity scores and hash_bytes from matches - Add .codegraph/ to gitignore (machine-local data, should not be committed) - Add CodeGraph configuration for multi-language indexing
471 lines
13 KiB
Python
471 lines
13 KiB
Python
# /// 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,
|
|
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,
|
|
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 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,
|
|
cfg_manager,
|
|
mo,
|
|
np,
|
|
pipeline,
|
|
room_nodes,
|
|
room_view_dataset,
|
|
save_object_image,
|
|
save_room_view,
|
|
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]
|
|
|
|
_text_labels = [
|
|
"a chair",
|
|
"a table",
|
|
"a sofa",
|
|
"a cabinet",
|
|
"a shelf",
|
|
"a lamp",
|
|
"a picture",
|
|
"a window",
|
|
"a door",
|
|
"a plant",
|
|
]
|
|
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}"
|
|
|
|
_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
|
|
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,
|
|
):
|
|
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")
|
|
|
|
_text_labels = [
|
|
"a chair",
|
|
"a table",
|
|
"a sofa",
|
|
"a cabinet",
|
|
"a shelf",
|
|
"a lamp",
|
|
"a picture",
|
|
"a window",
|
|
"a door",
|
|
"a plant",
|
|
]
|
|
_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,
|
|
}
|
|
|
|
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()
|