Mini-Nav/mini-nav/tests/test_image_utils.py

from unittest.mock import Mock

import torch
from PIL import Image

from utils.image import segment_image, segment_image_dataset


def test_segment_image_passes_pil_image_to_mask_generator() -> None:
    mock_generator = Mock(return_value={"masks": []})

    segment_image(
        mock_generator,
        Image.new("RGBA", (16, 16), color=(255, 0, 0, 255)),
        points_per_batch=32,
    )

    image_arg = mock_generator.call_args.args[0]
    assert isinstance(image_arg, Image.Image)
    assert image_arg.mode == "RGB"
    assert mock_generator.call_args.kwargs["points_per_batch"] == 32


def test_segment_image_supports_tensor_masks_output() -> None:
    masks_tensor = torch.tensor(
        [
            [
                [1, 1, 0],
                [1, 1, 0],
                [0, 0, 0],
            ],
            [
                [1, 1, 1],
                [1, 1, 1],
                [1, 1, 1],
            ],
        ],
        dtype=torch.float32,
    )
    mock_generator = Mock(return_value={"masks": masks_tensor})

    result = segment_image(
        mock_generator,
        Image.new("RGB", (3, 3), color=(0, 0, 0)),
        min_area=3,
        max_masks=5,
    )

    assert len(result) == 2
    assert result[0]["area"] == 9
    assert result[0]["bbox"] == [0, 0, 3, 3]
    assert result[1]["area"] == 4
    assert result[1]["bbox"] == [0, 0, 2, 2]


def test_segment_image_filters_tensor_masks_by_min_area() -> None:
    masks_tensor = torch.tensor(
        [
            [
                [1, 0, 0],
                [0, 0, 0],
                [0, 0, 0],
            ],
            [
                [1, 1, 0],
                [1, 1, 0],
                [0, 0, 0],
            ],
        ],
        dtype=torch.float32,
    )
    mock_generator = Mock(return_value={"masks": masks_tensor})

    result = segment_image(
        mock_generator,
        Image.new("RGB", (3, 3), color=(0, 0, 0)),
        min_area=2,
        max_masks=5,
    )

    assert len(result) == 1
    assert result[0]["area"] == 4


def test_segment_image_dataset_returns_per_image_masks_in_order() -> None:
    first_masks = torch.tensor(
        [[[1, 1, 0], [1, 1, 0], [0, 0, 0]]],
        dtype=torch.float32,
    )
    second_masks = torch.tensor(
        [[[1, 1, 1], [1, 1, 1], [1, 1, 1]]],
        dtype=torch.float32,
    )
    mock_generator = Mock(
        return_value=[{"masks": first_masks}, {"masks": second_masks}]
    )
    images = [
        Image.new("RGB", (3, 3), color=(0, 0, 0)),
        Image.new("RGB", (3, 3), color=(0, 0, 0)),
    ]

    result = segment_image_dataset(
        mock_generator,
        images,
        min_area=2,
        max_masks=5,
        points_per_batch=16,
    )

    assert len(result) == 2
    assert result[0][0]["area"] == 4
    assert result[1][0]["area"] == 9
    assert mock_generator.call_count == 1


def test_segment_image_dataset_falls_back_to_single_image_calls() -> None:
    call_index = {"value": 0}

    def fake_generator(images, points_per_batch):
        if isinstance(images, list):
            raise TypeError("Batch input unsupported")

        result_options = [
            {
                "masks": torch.tensor(
                    [[[1, 1, 0], [1, 1, 0], [0, 0, 0]]],
                    dtype=torch.float32,
                )
            },
            {
                "masks": torch.tensor(
                    [[[1, 1, 1], [1, 1, 1], [1, 1, 1]]],
                    dtype=torch.float32,
                )
            },
        ]
        out = result_options[call_index["value"]]
        call_index["value"] += 1
        return out

    images = [
        Image.new("RGB", (3, 3), color=(0, 0, 0)),
        Image.new("RGB", (3, 3), color=(0, 0, 0)),
    ]

    result = segment_image_dataset(
        fake_generator,
        images,
        min_area=2,
        max_masks=5,
        points_per_batch=16,
    )

    assert len(result) == 2
    assert result[0][0]["area"] == 4
    assert result[1][0]["area"] == 9