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:
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))]
|
||||
Reference in New Issue
Block a user