SikongJueluo/Mini-Nav
1
0
Fork 0
mirror of https://github.com/SikongJueluo/Mini-Nav.git synced 2026-07-12 20:15:31 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(compressors): add OWLv2 bbox crop to HashPipeline and refactor image utilities

Browse Source 
- Add Owlv2ForObjectDetection and Owlv2Processor imports to model_loader
- Refactor load_dino_model to return tuple of processor and model
- Rewrite generate_proposals_batch to group images by bbox count for efficient batching
- Add _normalize_single_bbox_list helper for bbox normalization
- Update verification.py to use new pipeline architecture with detect/segment/filter/crop steps
This commit is contained in:
SikongJueluo
2026-04-04 15:27:17 +08:00
parent 5f41cf5794
commit 94ed05a039
5 changed files with 679 additions and 586 deletions
Download Patch File Download Diff File Expand all files Collapse all files

458
notebooks/verification.py
View File

@@ -20,16 +20,20 @@ def _():
@app.cell
def global_base_deps():
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)
return Image, np, pl, torch
@app.cell
def project_setup():
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 (
@@ -39,8 +43,10 @@ def project_setup():
create_habitat_simulator,
render_topdown_scene_map,
)
from utils.image import numpy_to_pil
return (
HashPipeline,
HabitatSimulatorConfig,
ObjectNode,
RoomNode,
@@ -49,34 +55,26 @@ def project_setup():
cfg_manager,
collect_room_views_by_room,
create_habitat_simulator,
hamming_distance,
numpy_to_pil,
render_topdown_scene_map,
)
@app.cell
def setup_verification_context(
HabitatSimulatorConfig,
RoomNode,
create_habitat_simulator,
np,
):
scene_path = "data/scene_datasets/habitat-test-scenes/skokloster-castle.glb"
image_size = 512
num_rooms = 4
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
sam_max_masks = 5
sam_candidate_masks = 24
sam_min_area = 32 * 32
hash_bits = 512
pipeline_batch_size = 64
sim, agent = create_habitat_simulator(
HabitatSimulatorConfig(
scene_path=scene_path,
scene_path=_scene_path,
views_per_room=views_per_room,
image_size=image_size,
image_size=_image_size,
sensor_height=1.5,
move_forward_step=0.25,
enable_physics=False,
@@ -84,35 +82,25 @@ def setup_verification_context(
)
room_nodes = []
for idx in range(num_rooms):
point = sim.pathfinder.get_random_navigable_point()
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),
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(node.room_id, node.center)
return (
agent,
hash_bits,
meters_per_pixel,
pipeline_batch_size,
room_nodes,
sam_candidate_masks,
sam_max_masks,
sam_min_area,
sim,
views_per_room,
)
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_room_map(
def render_topdown(
TopDownSceneElements,
meters_per_pixel,
render_topdown_scene_map,
@@ -133,31 +121,25 @@ def render_topdown_room_map(
@app.cell
def build_scene_graph_pipeline(
Image,
ObjectNode,
SimpleSceneGraph,
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,
cfg_manager,
collect_room_views_by_room,
hash_bits,
np,
pipeline_batch_size,
numpy_to_pil,
room_nodes,
sam_candidate_masks,
sam_max_masks,
sam_min_area,
sim,
views_per_room,
):
from rich.progress import track
from compressors import MaskScoringConfig, rank_masks
from compressors.pipeline import HashPipeline
from compressors.proposal import extract_masked_region, generate_proposals_batch
from simulator import save_object_image, save_room_view
from utils import numpy_to_pil
all_room_views = collect_room_views_by_room(
agent=agent,
sim=sim,
@@ -165,188 +147,250 @@ def build_scene_graph_pipeline(
views_per_room=views_per_room,
)
hash_pipeline = HashPipeline(
dino_model="facebook/dinov2-large",
sam_model="facebook/sam2.1-hiera-large",
sam_min_mask_area=sam_min_area,
sam_max_masks=sam_max_masks,
hash_bits=hash_bits,
)
scene_graph = SimpleSceneGraph(
rooms={room.room_id: room for room in room_nodes},
objects={},
)
verification_output_dir = cfg_manager.get().output.directory / "verification"
verification_output_dir.mkdir(parents=True, exist_ok=True)
mask_scoring_config = MaskScoringConfig(
max_area_ratio=0.45,
reject_edge_touch_count=3,
reject_large_edge_touch_count=2,
reject_large_edge_area_ratio=0.12,
max_components=4,
min_largest_component_ratio=0.70,
)
total_masks = 0
total_raw_masks = 0
object_index = 0
# Flatten room views into (room_id, view_idx, PIL.Image) tuples.
room_view_dataset = [
(room_id, view_idx, rgb)
for room_id, views in all_room_views.items()
for view_idx, rgb in enumerate(views)
(_room_id, _view_idx, numpy_to_pil(_rgb))
for _room_id, _views in all_room_views.items()
for _view_idx, _rgb in enumerate(_views)
]
object_dataset = []
room_view_images = [numpy_to_pil(rgb) for _, _, rgb in room_view_dataset]
masks_dataset = generate_proposals_batch(
hash_pipeline.mask_generator,
room_view_images,
min_area=hash_pipeline.sam_min_mask_area,
max_masks=sam_candidate_masks,
points_per_batch=hash_pipeline.sam_points_per_batch,
)
if len(masks_dataset) != len(room_view_dataset):
raise RuntimeError("SAM dataset output size mismatch with room_view_dataset.")
dataset_jobs = list(zip(room_view_dataset, room_view_images, masks_dataset))
for (room_id, view_idx, _), image, masks in track(
dataset_jobs,
description="Building object dataset...",
):
save_room_view(verification_output_dir, room_id, view_idx, image)
total_raw_masks += len(masks)
ranked_masks = rank_masks(
masks,
image_shape=(image.height, image.width),
config=mask_scoring_config,
max_masks=sam_max_masks,
)
total_masks += len(ranked_masks)
for mask_idx, mask in enumerate(ranked_masks):
masked_image = extract_masked_region(image, mask["segment"])
object_dataset.append(
(room_id, view_idx, mask_idx, mask["bbox"], masked_image)
)
if object_dataset:
masked_images = [item[4] for item in object_dataset]
if any(not isinstance(img, Image.Image) for img in masked_images):
raise TypeError(
"object_dataset contains non-image entries for batch inference."
)
batched_bits = hash_pipeline.forward_dataset(
masked_images,
batch_size=pipeline_batch_size,
apply_sam=False,
)
if len(batched_bits) != len(object_dataset):
raise RuntimeError(
"Batch output size mismatch between masked images and hash outputs."
)
else:
batched_bits = []
for ob_idx, (room_id, view_idx, mask_idx, bbox, masked_image) in enumerate(
object_dataset
):
bits = batched_bits[ob_idx]
obj_center = np.array(
[bbox[0] + bbox[2] / 2, bbox[1] + bbox[3] / 2, 0.0],
dtype=np.float32,
)
obj_id = f"obj_{object_index:04d}"
object_index += 1
save_object_image(
verification_output_dir, room_id, obj_id, view_idx, mask_idx, masked_image
)
bits_array = np.asarray(bits.detach().cpu().numpy()).reshape(-1)
if bits_array.size == 512:
bits_binary = (bits_array > 0).astype(np.uint8)
hash_bytes = np.packbits(bits_binary).tobytes()
elif bits_array.size == 64:
hash_bytes = bits_array.astype(np.uint8).tobytes()
else:
raise ValueError(
f"Unexpected hash length: {bits_array.size}. Expected 512 bits or 64 bytes."
)
scene_graph.objects[obj_id] = ObjectNode(
obj_id=obj_id,
room_id=room_id,
position=obj_center,
visual_hash=hash_bytes,
semantic_hash=hash_bytes,
hit_count=1,
last_seen_frame=0,
)
print(f"Total objects created: {len(scene_graph.objects)}")
print(f"Total raw masks from SAM: {total_raw_masks}")
print(f"Total kept masks after ranking: {total_masks}")
print(f"Saved object images to: {verification_output_dir}")
return (scene_graph,)
print(f"Collected {len(room_view_dataset)} room views")
return all_room_views, room_view_dataset
@app.cell
def build_room_and_object_tables(pl, scene_graph):
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]),
"bbox_dx": float(room.bbox_extent[0]),
"bbox_dy": float(room.bbox_extent[1]),
"bbox_dz": float(room.bbox_extent[2]),
"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()
for _room in scene_graph.rooms.values()
]
object_rows = [
{
"obj_id": obj.obj_id,
"room_id": obj.room_id,
"last_seen_frame": int(obj.last_seen_frame),
"hit_count": int(obj.hit_count),
"visual_hash": obj.visual_hash.hex(),
"semantic_hash": obj.semantic_hash.hex(),
"visual_hash": obj.visual_hash.hex()[:16] + "...",
}
for obj in scene_graph.objects.values()
]
rooms_table = pl.DataFrame(room_rows)
objects_table = pl.DataFrame(object_rows)
return objects_table, rooms_table
rooms_df = pl.DataFrame(room_rows)
objects_df = pl.DataFrame(object_rows)
return objects_df, rooms_df
@app.cell
def upload_query_image(mo):
def upload_query(mo):
file_upload = mo.ui.file(
filetypes=["image/*"],
kind="area",
label="Upload a query image",
label="Upload a query image to find matching objects",
)
file_upload
return (file_upload,)
@app.cell
def _(file_upload, mo):
upload_image = None
if file_upload.value:
upload_image = mo.image(file_upload.contents(), alt="Uploaded query image")
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)
# Build a grid.
upload_image if upload_image is not None else mo.md("No image uploaded yet.")
return