feat(scenegraph): add ImageHashPipeline protocol and update query API

- Introduce ImageHashPipeline Protocol for extensible hash computation
- Rename query_crop_index to query_index for clarity
- Make query_crop nullable to handle missing crop edge case
- Add keyword-only arguments for better API clarity
- Handle empty scene graph objects gracefully
- Add comprehensive test coverage for query_image_against_scene_graph
This commit is contained in:
2026-05-21 15:03:39 +08:00
parent ba96cec406
commit 101a675ece
4 changed files with 371 additions and 30 deletions

View File

@@ -0,0 +1,276 @@
from __future__ import annotations
import sys
from dataclasses import dataclass
from pathlib import Path
import numpy as np
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
@dataclass(frozen=True)
class FakePipelineOutput:
hash_bits: torch.Tensor
cropped_images: list[Image.Image]
class FakePipeline:
def __init__(self, output: FakePipelineOutput):
self.output = output
self.calls: list[tuple[list[Image.Image], list[str], int]] = []
def process_batch(
self,
images: list[Image.Image],
text_labels: list[str],
batch_size: int = 32,
) -> FakePipelineOutput:
self.calls.append((images, text_labels, batch_size))
return self.output
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 _image(color: str = "white") -> Image.Image:
return Image.new("RGB", (8, 8), color=color)
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_single_query_hash_returns_expected_object_node():
query_bits = _bits_with_ones(1, 3)
query_hash = bits_tensor_to_hash_bytes(query_bits)
other_hash = _hash_with_ones(0)
query_crop = _image("blue")
pipeline = FakePipeline(
FakePipelineOutput(
hash_bits=query_bits.unsqueeze(0),
cropped_images=[query_crop],
)
)
graph = _scene_graph_with_hashes(("other", other_hash), ("target", query_hash))
image = _image("red")
results = query_image_against_scene_graph(
image,
pipeline=pipeline,
scene_graph=graph,
text_labels=["chair", "table"],
top_k=1,
batch_size=7,
)
assert len(results) == 1
assert pipeline.calls == [([image], ["chair", "table"], 7)]
assert results[0].query_index == 0
assert results[0].query_hash == query_hash
assert results[0].query_crop is query_crop
assert len(results[0].matches) == 1
assert results[0].matches[0].obj_id == "target"
assert results[0].matches[0].node is graph.objects["target"]
def test_one_image_with_multiple_query_hashes_returns_multiple_result_groups():
first_bits = _bits_with_ones(0)
second_bits = _bits_with_ones(2, 4)
hash_bits = torch.stack([first_bits, second_bits])
first_hash = bits_tensor_to_hash_bytes(first_bits)
second_hash = bits_tensor_to_hash_bytes(second_bits)
first_crop = _image("green")
second_crop = _image("yellow")
pipeline = FakePipeline(
FakePipelineOutput(
hash_bits=hash_bits,
cropped_images=[first_crop, second_crop],
)
)
graph = _scene_graph_with_hashes(("first", first_hash), ("second", second_hash))
results = query_image_against_scene_graph(
_image("red"),
pipeline=pipeline,
scene_graph=graph,
text_labels=["object"],
top_k=1,
)
assert [result.query_index for result in results] == [0, 1]
assert [result.query_hash for result in results] == [first_hash, second_hash]
assert [result.query_crop for result in results] == [first_crop, second_crop]
assert [result.matches[0].obj_id for result in results] == ["first", "second"]
def test_top_k_controls_match_count_and_preserves_cam_ordering():
query_bits = _bits_with_ones(0, 1, 2)
exact_hash = bits_tensor_to_hash_bytes(query_bits)
one_bit_off_hash = _hash_with_ones(0, 1)
two_bits_off_hash = _hash_with_ones(0)
pipeline = FakePipeline(
FakePipelineOutput(
hash_bits=query_bits.unsqueeze(0),
cropped_images=[_image("blue")],
)
)
graph = _scene_graph_with_hashes(
("two_bits_off", two_bits_off_hash),
("exact", exact_hash),
("one_bit_off", one_bit_off_hash),
)
results = query_image_against_scene_graph(
_image("red"),
pipeline=pipeline,
scene_graph=graph,
text_labels=["object"],
top_k=2,
)
assert [match.obj_id for match in results[0].matches] == ["exact", "one_bit_off"]
def test_empty_pipeline_output_returns_empty_list():
pipeline = FakePipeline(
FakePipelineOutput(
hash_bits=torch.empty((0, WIDTH), dtype=torch.int32),
cropped_images=[],
)
)
graph = _scene_graph_with_hashes(("obj", _hash_with_ones(0)))
results = query_image_against_scene_graph(
_image("red"),
pipeline=pipeline,
scene_graph=graph,
text_labels=["object"],
)
assert results == []
def test_empty_scene_graph_returns_result_group_with_empty_matches():
query_bits = _bits_with_ones(0)
query_crop = _image("blue")
pipeline = FakePipeline(
FakePipelineOutput(
hash_bits=query_bits.unsqueeze(0),
cropped_images=[query_crop],
)
)
results = query_image_against_scene_graph(
_image("red"),
pipeline=pipeline,
scene_graph=SimpleSceneGraph(),
text_labels=["object"],
)
assert len(results) == 1
assert results[0].query_index == 0
assert results[0].query_crop is query_crop
assert results[0].matches == []
def test_missing_crop_becomes_none_and_extra_crop_is_ignored():
first_bits = _bits_with_ones(0)
second_bits = _bits_with_ones(1)
first_hash = bits_tensor_to_hash_bytes(first_bits)
second_hash = bits_tensor_to_hash_bytes(second_bits)
only_crop = _image("blue")
pipeline = FakePipeline(
FakePipelineOutput(
hash_bits=torch.stack([first_bits, second_bits]),
cropped_images=[only_crop],
)
)
graph = _scene_graph_with_hashes(("first", first_hash), ("second", second_hash))
results = query_image_against_scene_graph(
_image("red"),
pipeline=pipeline,
scene_graph=graph,
text_labels=["object"],
)
assert [result.query_crop for result in results] == [only_crop, None]
extra_crop_pipeline = FakePipeline(
FakePipelineOutput(
hash_bits=first_bits.unsqueeze(0),
cropped_images=[only_crop, _image("yellow")],
)
)
extra_crop_results = query_image_against_scene_graph(
_image("red"),
pipeline=extra_crop_pipeline,
scene_graph=graph,
text_labels=["object"],
)
assert len(extra_crop_results) == 1
assert extra_crop_results[0].query_crop is only_crop
def test_scenegraph_package_exports_query_image_api():
from scenegraph import ( # noqa: PLC0415
ImageHashPipeline,
ImageSceneGraphQueryResult,
query_image_against_scene_graph as exported_query,
)
from scenegraph.query import ( # noqa: PLC0415
ImageHashPipeline as DirectImageHashPipeline,
ImageSceneGraphQueryResult as DirectImageSceneGraphQueryResult,
query_image_against_scene_graph as direct_query,
)
assert ImageHashPipeline is DirectImageHashPipeline
assert ImageSceneGraphQueryResult is DirectImageSceneGraphQueryResult
assert exported_query is direct_query
import scenegraph # noqa: PLC0415
required = {
"ImageHashPipeline",
"ImageSceneGraphQueryResult",
"query_image_against_scene_graph",
}
assert required <= set(scenegraph.__all__)