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

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