diff --git a/mini-nav/configs/config.yaml b/mini-nav/configs/config.yaml index 06840db..7079189 100644 --- a/mini-nav/configs/config.yaml +++ b/mini-nav/configs/config.yaml @@ -1,7 +1,7 @@ model: dino_model: "facebook/dinov2-large" compression_dim: 512 - device: "cuda:2" # auto-detect GPU + device: "cuda:3" # auto-detect GPU sam_model: "facebook/sam2.1-hiera-large" # SAM model name sam_min_mask_area: 100 # Minimum mask area threshold sam_max_masks: 10 # Maximum number of masks to keep diff --git a/notebooks/verification.py b/notebooks/verification.py index b1e6fe6..14b3bec 100644 --- a/notebooks/verification.py +++ b/notebooks/verification.py @@ -42,6 +42,8 @@ def project_imports(): 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 @@ -58,6 +60,8 @@ def project_imports(): hamming_distance, numpy_to_pil, render_topdown_scene_map, + save_object_image, + save_room_view, ) @@ -126,6 +130,8 @@ def pipeline_init(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,) @@ -167,6 +173,8 @@ def build_scene_graph( pipeline, room_nodes, room_view_dataset, + save_object_image, + save_room_view, ): scene_graph = SimpleSceneGraph( rooms={_room.room_id: _room for _room in room_nodes}, @@ -182,38 +190,77 @@ def build_scene_graph( _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) + _text_labels = [ + "a chair", + "a table", + "a sofa", + "a cabinet", + "a shelf", + "a lamp", + "a picture", + "a window", + "a door", + "a plant", + ] + _output = pipeline.process_batch( + _images, _text_labels, batch_size=32, return_debug_details=True + ) _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}" + from collections import Counter - _bits_array = _hash_bits.detach().cpu().numpy().reshape(-1) - _bits_binary = (_bits_array > 0).astype(np.uint8) - _hash_bytes = np.packbits(_bits_binary).tobytes() + _reasons = Counter(m["fallback_reason"] or "ok" for m in _output.debug_meta) + print(f"Fallback breakdown: {_reasons}") - object_images[_obj_id] = _cropped - _cropped.save(output_dir / f"{_obj_id}.png") + # Save original room views. + for _room_id, _view_idx, _image in room_view_dataset: + save_room_view(output_dir, _room_id, _view_idx, _image) - 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, - ) + # Prefix sum: map flat crop index to (input_image_idx, mask_idx). + _num_selected = [_m["num_selected"] for _m in _output.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 + for _img_idx, _n_crops in enumerate(_num_selected): + _room_id, _view_idx = _metadata[_img_idx] + for _mask_idx in range(_n_crops): + _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 _output.debug_meta if _meta["fallback_reason"] is not None ) - print(f"Created {len(scene_graph.objects)} objects") + print(f"Created {_obj_counter} objects") print(f"Saved cropped images to: {output_dir}") print(f"Fallback frames: {_fallback_count}/{len(_output.debug_meta)}") @@ -264,101 +311,116 @@ def query_matching( file_upload, hamming_distance, np, + mo, object_images, pipeline, scene_graph, torch, ): - import io + from io import BytesIO query_result = None query_cropped = None top_matches = [] - if file_upload.value: - _query_image = Image.open(io.BytesIO(file_upload.contents())).convert("RGB") + _file_contents = file_upload.contents() + mo.stop(not _file_contents, mo.md("请先上传文件")) - _text_labels = ["object"] - _output = pipeline.process_batch([_query_image], _text_labels, batch_size=1) - _query_bits = _output.hash_bits + _query_image = Image.open(BytesIO(_file_contents)).convert("RGB") - if _query_bits.numel() > 0: - query_cropped = _output.cropped_images[0] - _query_tensor = _query_bits[0].int() + _text_labels = [ + "a chair", + "a table", + "a sofa", + "a cabinet", + "a shelf", + "a lamp", + "a picture", + "a window", + "a door", + "a plant", + ] + _output = pipeline.process_batch([_query_image], _text_labels, batch_size=1) + _query_bits = (_output.hash_bits > 0).to(dtype=torch.int32) - _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 _query_bits.numel() > 0 and scene_graph.objects: + _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 + ].astype(np.int32) + _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) + _db_tensor = torch.tensor(np.array(_obj_hashes), dtype=torch.int32).to( + _query_bits.device + ) - _distances = hamming_distance(_query_tensor.unsqueeze(0), _db_tensor) - _distances = _distances.squeeze(0).cpu().numpy() + _distances = hamming_distance(_query_bits, _db_tensor) + _best_query_idx = int(_distances.min(dim=1).values.argmin().item()) - _top_k = min(5, len(_obj_ids)) - _top_indices = np.argsort(_distances)[:_top_k] + _query_tensor = _query_bits[_best_query_idx] + query_cropped = _output.cropped_images[_best_query_idx] + _query_distances = _distances[_best_query_idx].cpu().numpy() + _query_hash_hex = ( + np.packbits(_query_tensor.cpu().numpy().astype(np.uint8)).tobytes().hex() + ) - 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 - ] + _top_k = min(5, len(_obj_ids)) + _top_indices = np.argsort(_query_distances)[:_top_k] - query_result = { - "query_cropped": query_cropped, - "top_matches": top_matches, - } + top_matches = [ + { + "obj_id": _obj_ids[_i], + "distance": int(_query_distances[_i]), + "similarity": 1.0 - _query_distances[_i] / float(pipeline.hash_bits), + "hash_hex": scene_graph.objects[_obj_ids[_i]].visual_hash.hex(), + } + for _i in _top_indices + ] + + query_result = { + "query_cropped": query_cropped, + "query_hash_hex": _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): - if query_result is None: - mo.md("No query results yet. Upload an image above.") - else: - _result_items = [] + mo.stop(not query_result, mo.md("No query results yet. Upload an image above.")) - _result_items.append( - mo.vstack( - [ - mo.md("**Query (cropped)**"), - mo.image(query_cropped), - ], - align="center", - ) + _result_items = [ + 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) + for _match in top_matches: + _obj_id = _match["obj_id"] + _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", - ) + 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.hstack(_result_items, justify="center", gap=2) + mo.vstack(_result_items, justify="center", gap=2) return