Mini-Nav/tests/test_scenegraph_image_query.py

from __future__ import annotations

import sys
from pathlib import Path
from types import SimpleNamespace

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


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=pipeline, scene_graph=graph, text_labels=["a chair"], top_k=1, batch_size=7
    )

    assert len(results) == 1
    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
    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=pipeline, scene_graph=graph, text_labels=["a chair"], top_k=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"]


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=pipeline, scene_graph=graph, text_labels=["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=pipeline, scene_graph=SimpleSceneGraph(), text_labels=["object"]
    )

    assert results == []


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)
    )

    # 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 direct_query

    import scenegraph  # noqa: PLC0415

    required = {
        "ImageHashPipeline",
        "ImageSceneGraphQueryResult",
        "query_image_against_scene_graph",
    }
    assert required <= set(scenegraph.__all__)