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 .objectnode import ObjectNode
|
||||||
from .roomnode import RoomNode
|
from .roomnode import RoomNode
|
||||||
from .scenegraph import SimpleSceneGraph
|
from .scenegraph import SceneGraphMatch, SimpleSceneGraph
|
||||||
|
from .software_cam import CamMatch, SoftwareCamIndex, xnor_popcount_score
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
|
"CamMatch",
|
||||||
"ObjectNode",
|
"ObjectNode",
|
||||||
"RoomNode",
|
"RoomNode",
|
||||||
|
"SceneGraphMatch",
|
||||||
"SimpleSceneGraph",
|
"SimpleSceneGraph",
|
||||||
|
"SoftwareCamIndex",
|
||||||
"bits_tensor_to_cam_row",
|
"bits_tensor_to_cam_row",
|
||||||
"bits_tensor_to_hash_bytes",
|
"bits_tensor_to_hash_bytes",
|
||||||
"cam_row_to_hash_bytes",
|
"cam_row_to_hash_bytes",
|
||||||
"hash_bytes_to_bits_array",
|
"hash_bytes_to_bits_array",
|
||||||
"hash_bytes_to_cam_row",
|
"hash_bytes_to_cam_row",
|
||||||
|
"xnor_popcount_score",
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -1,10 +1,42 @@
|
|||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
|
|
||||||
from .roomnode import RoomNode
|
|
||||||
from .objectnode import ObjectNode
|
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
|
@dataclass
|
||||||
class SimpleSceneGraph:
|
class SimpleSceneGraph:
|
||||||
rooms: dict[str, RoomNode] = field(default_factory=dict)
|
rooms: dict[str, RoomNode] = field(default_factory=dict)
|
||||||
objects: dict[str, ObjectNode] = 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