diff --git a/mini-nav/scenegraph/__init__.py b/mini-nav/scenegraph/__init__.py index 8ba950d..b4ec631 100644 --- a/mini-nav/scenegraph/__init__.py +++ b/mini-nav/scenegraph/__init__.py @@ -2,7 +2,7 @@ This module exports the main scenegraph objects for easy import: -from mini_nav.scenegraph import SimpleSceneGraph, RoomNode, ObjectNode +from scenegraph import SimpleSceneGraph, RoomNode, ObjectNode """ from .hash_codec import ( @@ -13,13 +13,18 @@ from .hash_codec import ( hash_bytes_to_cam_row, ) from .objectnode import ObjectNode -from .query import ImageSceneGraphQueryResult, query_image_against_scene_graph +from .query import ( + ImageHashPipeline, + ImageSceneGraphQueryResult, + query_image_against_scene_graph, +) from .roomnode import RoomNode from .scenegraph import SceneGraphMatch, SimpleSceneGraph from .software_cam import CamMatch, SoftwareCamIndex, xnor_popcount_score __all__ = [ "CamMatch", + "ImageHashPipeline", "ImageSceneGraphQueryResult", "ObjectNode", "RoomNode", diff --git a/mini-nav/scenegraph/query.py b/mini-nav/scenegraph/query.py index ba57d9c..39c5621 100644 --- a/mini-nav/scenegraph/query.py +++ b/mini-nav/scenegraph/query.py @@ -1,7 +1,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Any +from typing import Any, Protocol, runtime_checkable from PIL import Image @@ -9,20 +9,40 @@ from .hash_codec import bits_tensor_to_hash_bytes from .scenegraph import SceneGraphMatch, SimpleSceneGraph +class ImageHashPipelineOutput(Protocol): + """Protocol for the output of ImageHashPipeline.process_batch.""" + + hash_bits: Any # torch.Tensor + cropped_images: list[Image.Image] + + +@runtime_checkable +class ImageHashPipeline(Protocol): + """Protocol for an image hash computation pipeline.""" + + def process_batch( + self, + images: list[Image.Image], + text_labels: list[str], + batch_size: int = 32, + ) -> ImageHashPipelineOutput: + ... + + @dataclass(frozen=True) class ImageSceneGraphQueryResult: - query_crop_index: int + query_index: int query_hash: bytes - query_crop: Image.Image + query_crop: Image.Image | None matches: list[SceneGraphMatch] def query_image_against_scene_graph( image: Image.Image, - pipeline: Any, + *, + pipeline: ImageHashPipeline, scene_graph: SimpleSceneGraph, text_labels: list[str], - *, top_k: int = 1, batch_size: int = 1, ) -> list[ImageSceneGraphQueryResult]: @@ -31,25 +51,24 @@ def query_image_against_scene_graph( cropped_images = list(output.cropped_images) if hash_bits.numel() == 0: - if cropped_images: - raise ValueError("hash_bits and cropped_images must align") return [] if hash_bits.dim() == 1: hash_bits = hash_bits.unsqueeze(0) - if hash_bits.shape[0] != len(cropped_images): - raise ValueError("hash_bits and cropped_images must align") - results: list[ImageSceneGraphQueryResult] = [] - for crop_index, query_bits in enumerate(hash_bits): + for query_index, query_bits in enumerate(hash_bits): query_hash = bits_tensor_to_hash_bytes(query_bits) - matches = scene_graph.query_by_visual_hash(query_hash, top_k=top_k) + if scene_graph.objects: + matches = scene_graph.query_by_visual_hash(query_hash, top_k=top_k) + else: + matches = [] + crop = cropped_images[query_index] if query_index < len(cropped_images) else None results.append( ImageSceneGraphQueryResult( - query_crop_index=crop_index, + query_index=query_index, query_hash=query_hash, - query_crop=cropped_images[crop_index], + query_crop=crop, matches=matches, ) ) diff --git a/tests/test_query_image_against_scene_graph.py b/tests/test_query_image_against_scene_graph.py new file mode 100644 index 0000000..85a6072 --- /dev/null +++ b/tests/test_query_image_against_scene_graph.py @@ -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__) diff --git a/tests/test_scenegraph_image_query.py b/tests/test_scenegraph_image_query.py index f047308..cff449c 100644 --- a/tests/test_scenegraph_image_query.py +++ b/tests/test_scenegraph_image_query.py @@ -5,7 +5,6 @@ from pathlib import Path from types import SimpleNamespace import numpy as np -import pytest import torch from PIL import Image @@ -88,11 +87,11 @@ def test_query_image_against_scene_graph_returns_exact_node_match(): pipeline = FakePipeline(query_bits.unsqueeze(0), [crop]) results = query_image_against_scene_graph( - _image(), pipeline, graph, ["a chair"], top_k=1, batch_size=7 + _image(), pipeline=pipeline, scene_graph=graph, text_labels=["a chair"], top_k=1, batch_size=7 ) assert len(results) == 1 - assert results[0].query_crop_index == 0 + assert results[0].query_index == 0 assert results[0].query_hash == query_hash assert results[0].query_crop is crop assert len(results[0].matches) == 1 @@ -114,10 +113,10 @@ def test_query_image_against_scene_graph_returns_one_result_per_query_crop(): 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 + _image(), pipeline=pipeline, scene_graph=graph, text_labels=["a chair"], top_k=1 ) - assert [result.query_crop_index for result in results] == [0, 1] + assert [result.query_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"] @@ -133,7 +132,7 @@ def test_query_image_against_scene_graph_preserves_topk_match_order(): pipeline = FakePipeline(query_bits.unsqueeze(0), [_image("red")]) results = query_image_against_scene_graph( - _image(), pipeline, graph, ["object"], top_k=3 + _image(), pipeline=pipeline, scene_graph=graph, text_labels=["object"], top_k=3 ) assert [match.obj_id for match in results[0].matches] == [ @@ -148,29 +147,71 @@ 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"] + _image(), pipeline=pipeline, scene_graph=SimpleSceneGraph(), text_labels=["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()]) +def test_missing_crop_returns_none_and_extra_crop_ignored(): + """Spec: missing crop index yields None, extra cropped images are 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") + graph = _scene_graph_with_hashes( + ("first", first_hash), ("second", second_hash) + ) - with pytest.raises(ValueError, match="hash_bits and cropped_images must align"): - query_image_against_scene_graph( - _image(), pipeline, SimpleSceneGraph(), ["object"] - ) + # Fewer cropped_images than hash_bits rows -> missing crop becomes None + missing_pipeline = FakePipeline( + torch.stack([first_bits, second_bits]), [only_crop] + ) + results = query_image_against_scene_graph( + _image("red"), + pipeline=missing_pipeline, + scene_graph=graph, + text_labels=["object"], + ) + assert [result.query_crop for result in results] == [only_crop, None] + + # More cropped_images than hash_bits rows -> extra crop is ignored + extra_pipeline = FakePipeline( + first_bits.unsqueeze(0), + [only_crop, _image("yellow")], + ) + extra_results = query_image_against_scene_graph( + _image("red"), + pipeline=extra_pipeline, + scene_graph=graph, + text_labels=["object"], + ) + assert len(extra_results) == 1 + assert extra_results[0].query_crop is only_crop def test_scenegraph_package_exports_image_query_api(): from scenegraph import ( # noqa: PLC0415 + ImageHashPipeline, ImageSceneGraphQueryResult, query_image_against_scene_graph as exported_query_image_against_scene_graph, ) 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_image_against_scene_graph is query_image_against_scene_graph + assert exported_query_image_against_scene_graph is direct_query + + import scenegraph # noqa: PLC0415 + + required = { + "ImageHashPipeline", + "ImageSceneGraphQueryResult", + "query_image_against_scene_graph", + } + assert required <= set(scenegraph.__all__)