Mini-Nav/tests/test_query_image_against_scene_graph.py

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