mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-13 04:25:32 +08:00
SikongJueluo
583b2156ea
- 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
85 lines
2.4 KiB
Python
85 lines
2.4 KiB
Python
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))]
|