refactor(verification): improve object creation / matching and image saving

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

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,22 +190,60 @@ 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
+
+    _reasons = Counter(m["fallback_reason"] or "ok" for m in _output.debug_meta)
+    print(f"Fallback breakdown: {_reasons}")
+
+    # Save original room views.
+    for _room_id, _view_idx, _image in room_view_dataset:
+        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 _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
-        _cropped.save(output_dir / f"{_obj_id}.png")
+            save_object_image(
+                output_dir, _room_id, _obj_id, _view_idx, _mask_idx, _cropped
+            )
 
             scene_graph.objects[_obj_id] = ObjectNode(
                 obj_id=_obj_id,
@@ -208,12 +254,13 @@ def build_scene_graph(
                 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,58 +311,78 @@ 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"]
+    _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",
+    ]
     _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()
+    _query_bits = (_output.hash_bits > 0).to(dtype=torch.int32)
 
+    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())
+
+        _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_k = min(5, len(_obj_ids))
-                _top_indices = np.argsort(_distances)[:_top_k]
+        _top_indices = np.argsort(_query_distances)[:_top_k]
 
         top_matches = [
             {
                 "obj_id": _obj_ids[_i],
-                        "distance": int(_distances[_i]),
-                        "similarity": 1.0 - _distances[_i] / float(pipeline.hash_bits),
+                "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,
         }
 
@@ -324,12 +391,9 @@ def query_matching(
 
 @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(
+    _result_items = [
         mo.vstack(
             [
                 mo.md("**Query (cropped)**"),
@@ -337,28 +401,26 @@ def display_results(mo, object_images, query_cropped, query_result, top_matches)
             ],
             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:
+        if _obj_img is not None:
             _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%}"),
+                        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