mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-13 04:25:32 +08:00
83 lines
2.0 KiB
Python
83 lines
2.0 KiB
Python
from unittest.mock import Mock
|
|
|
|
import torch
|
|
from PIL import Image
|
|
|
|
from utils.image import segment_image
|
|
|
|
|
|
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
|