mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-12 20:15:31 +08:00
feat(scenegraph): refactor image scene graph query into reusable function
- Export ImageSceneGraphQueryResult and query_image_against_scene_graph from scenegraph module - Replace inline hamming-distance-based image matching with dedicated query_image_against_scene_graph function - Improve top_matches structure by extracting similarity scores and hash_bytes from matches - Add .codegraph/ to gitignore (machine-local data, should not be committed) - Add CodeGraph configuration for multi-language indexing
This commit is contained in:
176
tests/test_scenegraph_image_query.py
Normal file
176
tests/test_scenegraph_image_query.py
Normal file
@@ -0,0 +1,176 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
|
||||
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 # noqa: E402
|
||||
from scenegraph.objectnode import ObjectNode # noqa: E402
|
||||
from scenegraph.query import query_image_against_scene_graph # noqa: E402
|
||||
from scenegraph.scenegraph import SimpleSceneGraph # noqa: E402
|
||||
|
||||
|
||||
WIDTH = 512
|
||||
|
||||
|
||||
class FakePipeline:
|
||||
def __init__(self, hash_bits: torch.Tensor, cropped_images: list[Image.Image]):
|
||||
self._hash_bits = hash_bits
|
||||
self._cropped_images = cropped_images
|
||||
self.calls = []
|
||||
|
||||
def process_batch(self, images, text_labels, batch_size=1):
|
||||
self.calls.append(
|
||||
{
|
||||
"images": images,
|
||||
"text_labels": text_labels,
|
||||
"batch_size": batch_size,
|
||||
}
|
||||
)
|
||||
return SimpleNamespace(
|
||||
hash_bits=self._hash_bits,
|
||||
cropped_images=self._cropped_images,
|
||||
debug_meta=[],
|
||||
)
|
||||
|
||||
|
||||
def _bits_with_ones(*indices: int) -> torch.Tensor:
|
||||
bits = torch.zeros(WIDTH, dtype=torch.int32)
|
||||
for index in indices:
|
||||
bits[index] = 1
|
||||
return bits
|
||||
|
||||
|
||||
def _hash_with_ones(*indices: int) -> bytes:
|
||||
return bits_tensor_to_hash_bytes(_bits_with_ones(*indices))
|
||||
|
||||
|
||||
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 _image(color: str = "white") -> Image.Image:
|
||||
return Image.new("RGB", (8, 8), color=color)
|
||||
|
||||
|
||||
def test_query_image_against_scene_graph_returns_exact_node_match():
|
||||
query_bits = _bits_with_ones(1, 2)
|
||||
query_hash = bits_tensor_to_hash_bytes(query_bits)
|
||||
graph = _scene_graph_with_hashes(
|
||||
("obj_a", _hash_with_ones(0)),
|
||||
("obj_b", query_hash),
|
||||
)
|
||||
crop = _image("red")
|
||||
pipeline = FakePipeline(query_bits.unsqueeze(0), [crop])
|
||||
|
||||
results = query_image_against_scene_graph(
|
||||
_image(), pipeline, graph, ["a chair"], top_k=1, batch_size=7
|
||||
)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].query_crop_index == 0
|
||||
assert results[0].query_hash == query_hash
|
||||
assert results[0].query_crop is crop
|
||||
assert len(results[0].matches) == 1
|
||||
assert results[0].matches[0].obj_id == "obj_b"
|
||||
assert results[0].matches[0].node is graph.objects["obj_b"]
|
||||
assert results[0].matches[0].score == WIDTH
|
||||
assert pipeline.calls[0]["text_labels"] == ["a chair"]
|
||||
assert pipeline.calls[0]["batch_size"] == 7
|
||||
|
||||
|
||||
def test_query_image_against_scene_graph_returns_one_result_per_query_crop():
|
||||
bits_a = _bits_with_ones(4)
|
||||
bits_b = _bits_with_ones(5, 6)
|
||||
hash_a = bits_tensor_to_hash_bytes(bits_a)
|
||||
hash_b = bits_tensor_to_hash_bytes(bits_b)
|
||||
graph = _scene_graph_with_hashes(("obj_a", hash_a), ("obj_b", hash_b))
|
||||
crop_a = _image("blue")
|
||||
crop_b = _image("green")
|
||||
pipeline = FakePipeline(torch.stack([bits_a, bits_b]), [crop_a, crop_b])
|
||||
|
||||
results = query_image_against_scene_graph(
|
||||
_image(), pipeline, graph, ["a chair"], top_k=1
|
||||
)
|
||||
|
||||
assert [result.query_crop_index for result in results] == [0, 1]
|
||||
assert [result.query_hash for result in results] == [hash_a, hash_b]
|
||||
assert [result.query_crop for result in results] == [crop_a, crop_b]
|
||||
assert [result.matches[0].obj_id for result in results] == ["obj_a", "obj_b"]
|
||||
|
||||
|
||||
def test_query_image_against_scene_graph_preserves_topk_match_order():
|
||||
query_bits = _bits_with_ones(0, 1, 2)
|
||||
graph = _scene_graph_with_hashes(
|
||||
("obj_far", _hash_with_ones(0)),
|
||||
("obj_exact", bits_tensor_to_hash_bytes(query_bits)),
|
||||
("obj_near", _hash_with_ones(0, 1)),
|
||||
)
|
||||
pipeline = FakePipeline(query_bits.unsqueeze(0), [_image("red")])
|
||||
|
||||
results = query_image_against_scene_graph(
|
||||
_image(), pipeline, graph, ["object"], top_k=3
|
||||
)
|
||||
|
||||
assert [match.obj_id for match in results[0].matches] == [
|
||||
"obj_exact",
|
||||
"obj_near",
|
||||
"obj_far",
|
||||
]
|
||||
assert [match.score for match in results[0].matches] == [WIDTH, WIDTH - 1, WIDTH - 2]
|
||||
|
||||
|
||||
def test_query_image_against_scene_graph_returns_empty_list_for_no_hashes():
|
||||
pipeline = FakePipeline(torch.empty((0, WIDTH), dtype=torch.int32), [])
|
||||
|
||||
results = query_image_against_scene_graph(
|
||||
_image(), pipeline, SimpleSceneGraph(), ["object"]
|
||||
)
|
||||
|
||||
assert results == []
|
||||
|
||||
|
||||
def test_query_image_against_scene_graph_rejects_hash_crop_count_mismatch():
|
||||
pipeline = FakePipeline(torch.stack([_bits_with_ones(0), _bits_with_ones(1)]), [_image()])
|
||||
|
||||
with pytest.raises(ValueError, match="hash_bits and cropped_images must align"):
|
||||
query_image_against_scene_graph(
|
||||
_image(), pipeline, SimpleSceneGraph(), ["object"]
|
||||
)
|
||||
|
||||
|
||||
def test_scenegraph_package_exports_image_query_api():
|
||||
from scenegraph import ( # noqa: PLC0415
|
||||
ImageSceneGraphQueryResult,
|
||||
query_image_against_scene_graph as exported_query_image_against_scene_graph,
|
||||
)
|
||||
from scenegraph.query import ( # noqa: PLC0415
|
||||
ImageSceneGraphQueryResult as DirectImageSceneGraphQueryResult,
|
||||
)
|
||||
|
||||
assert ImageSceneGraphQueryResult is DirectImageSceneGraphQueryResult
|
||||
assert exported_query_image_against_scene_graph is query_image_against_scene_graph
|
||||
Reference in New Issue
Block a user