diff --git a/.codegraph/.gitignore b/.codegraph/.gitignore new file mode 100644 index 0000000..9de0f16 --- /dev/null +++ b/.codegraph/.gitignore @@ -0,0 +1,16 @@ +# CodeGraph data files +# These are local to each machine and should not be committed + +# Database +*.db +*.db-wal +*.db-shm + +# Cache +cache/ + +# Logs +*.log + +# Hook markers +.dirty diff --git a/.codegraph/config.json b/.codegraph/config.json new file mode 100644 index 0000000..7af60ad --- /dev/null +++ b/.codegraph/config.json @@ -0,0 +1,143 @@ +{ + "version": 1, + "include": [ + "**/*.ts", + "**/*.tsx", + "**/*.js", + "**/*.jsx", + "**/*.py", + "**/*.go", + "**/*.rs", + "**/*.java", + "**/*.c", + "**/*.h", + "**/*.cpp", + "**/*.hpp", + "**/*.cc", + "**/*.cxx", + "**/*.cs", + "**/*.php", + "**/*.rb", + "**/*.swift", + "**/*.kt", + "**/*.kts", + "**/*.dart", + "**/*.svelte", + "**/*.vue", + "**/*.liquid", + "**/*.pas", + "**/*.dpr", + "**/*.dpk", + "**/*.lpr", + "**/*.dfm", + "**/*.fmx", + "**/*.scala", + "**/*.sc" + ], + "exclude": [ + "**/.git/**", + "**/node_modules/**", + "**/vendor/**", + "**/Pods/**", + "**/dist/**", + "**/build/**", + "**/out/**", + "**/bin/**", + "**/obj/**", + "**/target/**", + "**/*.min.js", + "**/*.bundle.js", + "**/.next/**", + "**/.nuxt/**", + "**/.svelte-kit/**", + "**/.output/**", + "**/.turbo/**", + "**/.cache/**", + "**/.parcel-cache/**", + "**/.vite/**", + "**/.astro/**", + "**/.docusaurus/**", + "**/.gatsby/**", + "**/.webpack/**", + "**/.nx/**", + "**/.yarn/cache/**", + "**/.pnpm-store/**", + "**/storybook-static/**", + "**/.expo/**", + "**/web-build/**", + "**/ios/Pods/**", + "**/ios/build/**", + "**/android/build/**", + "**/android/.gradle/**", + "**/__pycache__/**", + "**/.venv/**", + "**/venv/**", + "**/site-packages/**", + "**/dist-packages/**", + "**/.pytest_cache/**", + "**/.mypy_cache/**", + "**/.ruff_cache/**", + "**/.tox/**", + "**/.nox/**", + "**/*.egg-info/**", + "**/.eggs/**", + "**/go/pkg/mod/**", + "**/target/debug/**", + "**/target/release/**", + "**/.gradle/**", + "**/.m2/**", + "**/generated-sources/**", + "**/.kotlin/**", + "**/.dart_tool/**", + "**/.vs/**", + "**/.nuget/**", + "**/artifacts/**", + "**/publish/**", + "**/cmake-build-*/**", + "**/CMakeFiles/**", + "**/bazel-*/**", + "**/vcpkg_installed/**", + "**/.conan/**", + "**/Debug/**", + "**/Release/**", + "**/x64/**", + "**/.pio/**", + "**/release/**", + "**/*.app/**", + "**/*.asar", + "**/DerivedData/**", + "**/.build/**", + "**/.swiftpm/**", + "**/xcuserdata/**", + "**/Carthage/Build/**", + "**/SourcePackages/**", + "**/__history/**", + "**/__recovery/**", + "**/*.dcu", + "**/.composer/**", + "**/storage/framework/**", + "**/bootstrap/cache/**", + "**/.bundle/**", + "**/tmp/cache/**", + "**/public/assets/**", + "**/public/packs/**", + "**/.yardoc/**", + "**/coverage/**", + "**/htmlcov/**", + "**/.nyc_output/**", + "**/test-results/**", + "**/.coverage/**", + "**/.idea/**", + "**/logs/**", + "**/tmp/**", + "**/temp/**", + "**/_build/**", + "**/docs/_build/**", + "**/site/**" + ], + "languages": [], + "frameworks": [], + "maxFileSize": 1048576, + "extractDocstrings": true, + "trackCallSites": true +} \ No newline at end of file diff --git a/mini-nav/scenegraph/__init__.py b/mini-nav/scenegraph/__init__.py index 3b3caab..8ba950d 100644 --- a/mini-nav/scenegraph/__init__.py +++ b/mini-nav/scenegraph/__init__.py @@ -13,12 +13,14 @@ from .hash_codec import ( hash_bytes_to_cam_row, ) from .objectnode import ObjectNode +from .query import 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", + "ImageSceneGraphQueryResult", "ObjectNode", "RoomNode", "SceneGraphMatch", @@ -29,5 +31,6 @@ __all__ = [ "cam_row_to_hash_bytes", "hash_bytes_to_bits_array", "hash_bytes_to_cam_row", + "query_image_against_scene_graph", "xnor_popcount_score", ] diff --git a/mini-nav/scenegraph/query.py b/mini-nav/scenegraph/query.py new file mode 100644 index 0000000..ba57d9c --- /dev/null +++ b/mini-nav/scenegraph/query.py @@ -0,0 +1,57 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from PIL import Image + +from .hash_codec import bits_tensor_to_hash_bytes +from .scenegraph import SceneGraphMatch, SimpleSceneGraph + + +@dataclass(frozen=True) +class ImageSceneGraphQueryResult: + query_crop_index: int + query_hash: bytes + query_crop: Image.Image + matches: list[SceneGraphMatch] + + +def query_image_against_scene_graph( + image: Image.Image, + pipeline: Any, + scene_graph: SimpleSceneGraph, + text_labels: list[str], + *, + top_k: int = 1, + batch_size: int = 1, +) -> list[ImageSceneGraphQueryResult]: + output = pipeline.process_batch([image], text_labels, batch_size=batch_size) + hash_bits = output.hash_bits + 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): + query_hash = bits_tensor_to_hash_bytes(query_bits) + matches = scene_graph.query_by_visual_hash(query_hash, top_k=top_k) + results.append( + ImageSceneGraphQueryResult( + query_crop_index=crop_index, + query_hash=query_hash, + query_crop=cropped_images[crop_index], + matches=matches, + ) + ) + + return results diff --git a/notebooks/verification.py b/notebooks/verification.py index 516b6ae..0fef577 100644 --- a/notebooks/verification.py +++ b/notebooks/verification.py @@ -33,9 +33,14 @@ def base_dependencies(): @app.cell def project_imports(): """Project module imports using new architecture.""" - from compressors import HashPipeline, hamming_distance + from compressors import HashPipeline from configs import cfg_manager - from scenegraph import ObjectNode, RoomNode, SimpleSceneGraph + from scenegraph import ( + ObjectNode, + RoomNode, + SimpleSceneGraph, + query_image_against_scene_graph, + ) from simulator import ( HabitatSimulatorConfig, TopDownSceneElements, @@ -57,8 +62,8 @@ def project_imports(): cfg_manager, collect_room_views_by_room, create_habitat_simulator, - hamming_distance, numpy_to_pil, + query_image_against_scene_graph, render_topdown_scene_map, save_object_image, save_room_view, @@ -362,13 +367,11 @@ def upload_query(mo): def query_matching( Image, file_upload, - hamming_distance, - np, mo, object_images, pipeline, + query_image_against_scene_graph, scene_graph, - torch, ): from io import BytesIO @@ -393,49 +396,34 @@ def query_matching( "a door", "a plant", ] - _output = pipeline.process_batch([_query_image], _text_labels, batch_size=1) - _query_bits = (_output.hash_bits > 0).to(dtype=torch.int32) + _query_results = query_image_against_scene_graph( + image=_query_image, + pipeline=pipeline, + scene_graph=scene_graph, + text_labels=_text_labels, + top_k=5, + batch_size=1, + ) - if _query_bits.numel() > 0 and scene_graph.objects: - _obj_ids = list(scene_graph.objects.keys()) - _obj_hashes = [] - for _obj_id in _obj_ids: - _obj = scene_graph.objects[_obj_id] - _bits = np.unpackbits(np.frombuffer(_obj.visual_hash, dtype=np.uint8))[ - : pipeline.hash_bits - ].astype(np.int32) - _obj_hashes.append(_bits) - - _db_tensor = torch.tensor(np.array(_obj_hashes), dtype=torch.int32).to( - _query_bits.device + if _query_results: + _best_result = max( + _query_results, + key=lambda result: result.matches[0].score if result.matches else -1, ) - - _distances = hamming_distance(_query_bits, _db_tensor) - _best_query_idx = int(_distances.min(dim=1).values.argmin().item()) - - _query_tensor = _query_bits[_best_query_idx] - query_cropped = _output.cropped_images[_best_query_idx] - _query_distances = _distances[_best_query_idx].cpu().numpy() - _query_hash_hex = ( - np.packbits(_query_tensor.cpu().numpy().astype(np.uint8)).tobytes().hex() - ) - - _top_k = min(5, len(_obj_ids)) - _top_indices = np.argsort(_query_distances)[:_top_k] - + query_cropped = _best_result.query_crop top_matches = [ { - "obj_id": _obj_ids[_i], - "distance": int(_query_distances[_i]), - "similarity": 1.0 - _query_distances[_i] / float(pipeline.hash_bits), - "hash_hex": scene_graph.objects[_obj_ids[_i]].visual_hash.hex(), + "obj_id": match.obj_id, + "distance": int(pipeline.hash_bits - match.score), + "similarity": match.similarity, + "hash_hex": match.hash_bytes.hex(), } - for _i in _top_indices + for match in _best_result.matches ] query_result = { "query_cropped": query_cropped, - "query_hash_hex": _query_hash_hex, + "query_hash_hex": _best_result.query_hash.hex(), "top_matches": top_matches, } diff --git a/tests/test_scenegraph_image_query.py b/tests/test_scenegraph_image_query.py new file mode 100644 index 0000000..f047308 --- /dev/null +++ b/tests/test_scenegraph_image_query.py @@ -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