mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-12 20:15:31 +08:00
116 lines
2.8 KiB
Python
116 lines
2.8 KiB
Python
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 = {
|
|
"masks": torch.tensor(
|
|
[[[1, 1, 0], [1, 1, 0], [0, 0, 0]]],
|
|
dtype=torch.float32,
|
|
)
|
|
}
|
|
second_masks = {
|
|
"masks": torch.tensor(
|
|
[[[1, 1, 1], [1, 1, 1], [1, 1, 1]]],
|
|
dtype=torch.float32,
|
|
)
|
|
}
|
|
mock_generator = Mock(side_effect=[first_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 == 2
|