mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-12 20:15:31 +08:00
feat(scenegraph): add SoftwareCamIndex for visual hash similarity queries
- Add SoftwareCamIndex class with xnor_popcount_score for CAM-style matching - Add CamMatch and SceneGraphMatch dataclasses for query results - Add query_by_visual_hash method to SimpleSceneGraph - Add comprehensive tests for SoftwareCamIndex and xnor_popcount_score
This commit is contained in:
@@ -14,15 +14,20 @@ from .hash_codec import (
|
||||
)
|
||||
from .objectnode import ObjectNode
|
||||
from .roomnode import RoomNode
|
||||
from .scenegraph import SimpleSceneGraph
|
||||
from .scenegraph import SceneGraphMatch, SimpleSceneGraph
|
||||
from .software_cam import CamMatch, SoftwareCamIndex, xnor_popcount_score
|
||||
|
||||
__all__ = [
|
||||
"CamMatch",
|
||||
"ObjectNode",
|
||||
"RoomNode",
|
||||
"SceneGraphMatch",
|
||||
"SimpleSceneGraph",
|
||||
"SoftwareCamIndex",
|
||||
"bits_tensor_to_cam_row",
|
||||
"bits_tensor_to_hash_bytes",
|
||||
"cam_row_to_hash_bytes",
|
||||
"hash_bytes_to_bits_array",
|
||||
"hash_bytes_to_cam_row",
|
||||
"xnor_popcount_score",
|
||||
]
|
||||
|
||||
@@ -1,10 +1,42 @@
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from .roomnode import RoomNode
|
||||
from .objectnode import ObjectNode
|
||||
from .roomnode import RoomNode
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SceneGraphMatch:
|
||||
obj_id: str
|
||||
node: ObjectNode
|
||||
row_index: int
|
||||
score: int
|
||||
similarity: float
|
||||
hash_bytes: bytes
|
||||
|
||||
|
||||
@dataclass
|
||||
class SimpleSceneGraph:
|
||||
rooms: dict[str, RoomNode] = field(default_factory=dict)
|
||||
objects: dict[str, ObjectNode] = field(default_factory=dict)
|
||||
|
||||
def query_by_visual_hash(
|
||||
self,
|
||||
query_hash_bytes: bytes,
|
||||
*,
|
||||
top_k: int = 1,
|
||||
) -> list[SceneGraphMatch]:
|
||||
from .software_cam import SoftwareCamIndex
|
||||
|
||||
index = SoftwareCamIndex.from_scene_graph(self)
|
||||
cam_matches = index.query(query_hash_bytes, top_k=top_k)
|
||||
return [
|
||||
SceneGraphMatch(
|
||||
obj_id=match.obj_id,
|
||||
node=self.objects[match.obj_id],
|
||||
row_index=match.row_index,
|
||||
score=match.score,
|
||||
similarity=match.similarity,
|
||||
hash_bytes=match.hash_bytes,
|
||||
)
|
||||
for match in cam_matches
|
||||
]
|
||||
|
||||
84
mini-nav/scenegraph/software_cam.py
Normal file
84
mini-nav/scenegraph/software_cam.py
Normal file
@@ -0,0 +1,84 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from .hash_codec import DEFAULT_HASH_WIDTH, hash_bytes_to_cam_row
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .scenegraph import SimpleSceneGraph
|
||||
|
||||
|
||||
def xnor_popcount_score(
|
||||
query_row: int,
|
||||
stored_row: int,
|
||||
*,
|
||||
width: int = DEFAULT_HASH_WIDTH,
|
||||
) -> int:
|
||||
"""Compute CAM-style same-bit score for two integer hash rows."""
|
||||
if width <= 0:
|
||||
raise ValueError("width must be greater than 0")
|
||||
mask = (1 << width) - 1
|
||||
return int((~(int(query_row) ^ int(stored_row)) & mask).bit_count())
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CamMatch:
|
||||
row_index: int
|
||||
obj_id: str
|
||||
score: int
|
||||
similarity: float
|
||||
hash_bytes: bytes
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SoftwareCamIndex:
|
||||
obj_ids: tuple[str, ...]
|
||||
rows: tuple[int, ...]
|
||||
hashes: tuple[bytes, ...]
|
||||
width: int = DEFAULT_HASH_WIDTH
|
||||
|
||||
@classmethod
|
||||
def from_scene_graph(
|
||||
cls,
|
||||
scene_graph: "SimpleSceneGraph",
|
||||
*,
|
||||
width: int = DEFAULT_HASH_WIDTH,
|
||||
) -> "SoftwareCamIndex":
|
||||
obj_ids: list[str] = []
|
||||
rows: list[int] = []
|
||||
hashes: list[bytes] = []
|
||||
|
||||
for obj_id, node in scene_graph.objects.items():
|
||||
hash_bytes = node.visual_hash
|
||||
obj_ids.append(obj_id)
|
||||
hashes.append(hash_bytes)
|
||||
rows.append(hash_bytes_to_cam_row(hash_bytes, width=width))
|
||||
|
||||
return cls(
|
||||
obj_ids=tuple(obj_ids),
|
||||
rows=tuple(rows),
|
||||
hashes=tuple(hashes),
|
||||
width=width,
|
||||
)
|
||||
|
||||
def query(self, query_hash_bytes: bytes, *, top_k: int = 1) -> list[CamMatch]:
|
||||
if top_k <= 0:
|
||||
raise ValueError("top_k must be greater than 0")
|
||||
if not self.rows:
|
||||
raise ValueError("cannot query an empty SoftwareCamIndex")
|
||||
|
||||
query_row = hash_bytes_to_cam_row(query_hash_bytes, width=self.width)
|
||||
matches = [
|
||||
CamMatch(
|
||||
row_index=row_index,
|
||||
obj_id=self.obj_ids[row_index],
|
||||
score=score,
|
||||
similarity=score / float(self.width),
|
||||
hash_bytes=self.hashes[row_index],
|
||||
)
|
||||
for row_index, row in enumerate(self.rows)
|
||||
for score in [xnor_popcount_score(query_row, row, width=self.width)]
|
||||
]
|
||||
matches.sort(key=lambda match: (-match.score, match.row_index))
|
||||
return matches[: min(top_k, len(matches))]
|
||||
191
tests/test_software_cam_index.py
Normal file
191
tests/test_software_cam_index.py
Normal file
@@ -0,0 +1,191 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
|
||||
MINI_NAV_DIR = Path(__file__).resolve().parents[1] / "mini-nav"
|
||||
sys.path.insert(0, str(MINI_NAV_DIR))
|
||||
|
||||
from scenegraph.hash_codec import bits_tensor_to_hash_bytes, hash_bytes_to_cam_row # noqa: E402
|
||||
from scenegraph.objectnode import ObjectNode # noqa: E402
|
||||
from scenegraph.scenegraph import SimpleSceneGraph # noqa: E402
|
||||
from scenegraph.software_cam import SoftwareCamIndex, xnor_popcount_score # noqa: E402
|
||||
|
||||
|
||||
WIDTH = 512
|
||||
|
||||
|
||||
def _hash_with_ones(*indices: int) -> bytes:
|
||||
bits = torch.zeros(WIDTH, dtype=torch.int32)
|
||||
for index in indices:
|
||||
bits[index] = 1
|
||||
return bits_tensor_to_hash_bytes(bits)
|
||||
|
||||
|
||||
def _node(obj_id: str, hash_bytes: bytes) -> ObjectNode:
|
||||
return ObjectNode(
|
||||
obj_id=obj_id,
|
||||
room_id="room_a",
|
||||
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=0,
|
||||
)
|
||||
|
||||
|
||||
def _scene_graph_with_hashes(*items: tuple[str, bytes]) -> SimpleSceneGraph:
|
||||
graph = SimpleSceneGraph()
|
||||
for obj_id, hash_bytes in items:
|
||||
graph.objects[obj_id] = _node(obj_id, hash_bytes)
|
||||
return graph
|
||||
|
||||
|
||||
def test_xnor_popcount_score_counts_matching_bits():
|
||||
assert xnor_popcount_score(0b00, 0b00, width=2) == 2
|
||||
assert xnor_popcount_score(0b00, 0b01, width=2) == 1
|
||||
assert xnor_popcount_score(0b00, 0b11, width=2) == 0
|
||||
|
||||
|
||||
def test_software_cam_index_from_scene_graph_preserves_snapshot_order_and_rows():
|
||||
hash_a = _hash_with_ones(0)
|
||||
hash_b = _hash_with_ones(1)
|
||||
graph = _scene_graph_with_hashes(("obj_a", hash_a), ("obj_b", hash_b))
|
||||
|
||||
index = SoftwareCamIndex.from_scene_graph(graph)
|
||||
|
||||
assert index.obj_ids == ("obj_a", "obj_b")
|
||||
assert index.hashes == (hash_a, hash_b)
|
||||
assert index.rows == (
|
||||
hash_bytes_to_cam_row(hash_a),
|
||||
hash_bytes_to_cam_row(hash_b),
|
||||
)
|
||||
assert index.width == WIDTH
|
||||
|
||||
|
||||
def test_software_cam_index_query_returns_top1_exact_match():
|
||||
hash_a = _hash_with_ones(0)
|
||||
hash_b = _hash_with_ones(3, 5)
|
||||
graph = _scene_graph_with_hashes(("obj_a", hash_a), ("obj_b", hash_b))
|
||||
index = SoftwareCamIndex.from_scene_graph(graph)
|
||||
|
||||
matches = index.query(hash_b)
|
||||
|
||||
assert len(matches) == 1
|
||||
assert matches[0].obj_id == "obj_b"
|
||||
assert matches[0].row_index == 1
|
||||
assert matches[0].score == WIDTH
|
||||
assert matches[0].similarity == 1.0
|
||||
assert matches[0].hash_bytes == hash_b
|
||||
|
||||
|
||||
def test_software_cam_index_query_topk_orders_by_score_descending():
|
||||
query = _hash_with_ones(0, 1, 2)
|
||||
two_bits_different = _hash_with_ones(0)
|
||||
exact = _hash_with_ones(0, 1, 2)
|
||||
one_bit_different = _hash_with_ones(0, 1)
|
||||
graph = _scene_graph_with_hashes(
|
||||
("obj_a", two_bits_different),
|
||||
("obj_b", exact),
|
||||
("obj_c", one_bit_different),
|
||||
)
|
||||
index = SoftwareCamIndex.from_scene_graph(graph)
|
||||
|
||||
matches = index.query(query, top_k=3)
|
||||
|
||||
assert [match.obj_id for match in matches] == ["obj_b", "obj_c", "obj_a"]
|
||||
assert [match.score for match in matches] == [WIDTH, WIDTH - 1, WIDTH - 2]
|
||||
|
||||
|
||||
def test_software_cam_index_query_tiebreaks_by_smaller_row_index():
|
||||
query = _hash_with_ones(0)
|
||||
equal_a = _hash_with_ones(1)
|
||||
equal_b = _hash_with_ones(2)
|
||||
graph = _scene_graph_with_hashes(("obj_a", equal_a), ("obj_b", equal_b))
|
||||
index = SoftwareCamIndex.from_scene_graph(graph)
|
||||
|
||||
matches = index.query(query, top_k=2)
|
||||
|
||||
assert [match.obj_id for match in matches] == ["obj_a", "obj_b"]
|
||||
assert matches[0].score == matches[1].score
|
||||
assert [match.row_index for match in matches] == [0, 1]
|
||||
|
||||
|
||||
def test_software_cam_index_query_topk_larger_than_rows_returns_all_rows():
|
||||
hash_a = _hash_with_ones(0)
|
||||
hash_b = _hash_with_ones(1)
|
||||
graph = _scene_graph_with_hashes(("obj_a", hash_a), ("obj_b", hash_b))
|
||||
index = SoftwareCamIndex.from_scene_graph(graph)
|
||||
|
||||
matches = index.query(hash_a, top_k=99)
|
||||
|
||||
assert len(matches) == 2
|
||||
|
||||
|
||||
def test_empty_software_cam_index_can_be_built_but_not_queried():
|
||||
index = SoftwareCamIndex.from_scene_graph(SimpleSceneGraph())
|
||||
query = _hash_with_ones(0)
|
||||
|
||||
assert index.obj_ids == ()
|
||||
assert index.rows == ()
|
||||
assert index.hashes == ()
|
||||
with pytest.raises(ValueError, match="empty SoftwareCamIndex"):
|
||||
index.query(query)
|
||||
|
||||
|
||||
def test_software_cam_index_query_rejects_invalid_top_k():
|
||||
hash_a = _hash_with_ones(0)
|
||||
graph = _scene_graph_with_hashes(("obj_a", hash_a))
|
||||
index = SoftwareCamIndex.from_scene_graph(graph)
|
||||
|
||||
with pytest.raises(ValueError, match="top_k"):
|
||||
index.query(hash_a, top_k=0)
|
||||
|
||||
|
||||
def test_software_cam_index_uses_snapshot_semantics():
|
||||
hash_a = _hash_with_ones(0)
|
||||
hash_b = _hash_with_ones(1)
|
||||
graph = _scene_graph_with_hashes(("obj_a", hash_a))
|
||||
|
||||
index = SoftwareCamIndex.from_scene_graph(graph)
|
||||
graph.objects["obj_b"] = _node("obj_b", hash_b)
|
||||
|
||||
assert index.obj_ids == ("obj_a",)
|
||||
assert "obj_b" not in index.obj_ids
|
||||
|
||||
|
||||
def test_simple_scene_graph_query_by_visual_hash_returns_scene_graph_match():
|
||||
hash_a = _hash_with_ones(0)
|
||||
hash_b = _hash_with_ones(4, 8)
|
||||
graph = _scene_graph_with_hashes(("obj_a", hash_a), ("obj_b", hash_b))
|
||||
|
||||
matches = graph.query_by_visual_hash(hash_b, top_k=1)
|
||||
|
||||
assert len(matches) == 1
|
||||
assert matches[0].obj_id == "obj_b"
|
||||
assert matches[0].node is graph.objects["obj_b"]
|
||||
assert matches[0].row_index == 1
|
||||
assert matches[0].score == WIDTH
|
||||
assert matches[0].similarity == 1.0
|
||||
assert matches[0].hash_bytes == hash_b
|
||||
|
||||
|
||||
def test_scenegraph_package_exports_software_cam_query_api():
|
||||
from scenegraph import ( # noqa: PLC0415
|
||||
CamMatch,
|
||||
SceneGraphMatch,
|
||||
SoftwareCamIndex as ExportedSoftwareCamIndex,
|
||||
xnor_popcount_score as exported_xnor_popcount_score,
|
||||
)
|
||||
from scenegraph.scenegraph import SceneGraphMatch as DirectSceneGraphMatch # noqa: PLC0415
|
||||
from scenegraph.software_cam import CamMatch as DirectCamMatch # noqa: PLC0415
|
||||
|
||||
assert CamMatch is DirectCamMatch
|
||||
assert SceneGraphMatch is DirectSceneGraphMatch
|
||||
assert ExportedSoftwareCamIndex is SoftwareCamIndex
|
||||
assert exported_xnor_popcount_score is xnor_popcount_score
|
||||
Reference in New Issue
Block a user