195 lines
6.4 KiB
Python
195 lines
6.4 KiB
Python
|
# Copyright (C) CVAT.ai Corporation
|
||
|
|
#
|
||
|
|
# SPDX-License-Identifier: MIT
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import textwrap
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
try:
|
||
|
|
import numpy as np
|
||
|
|
from cvat_sdk.masks import decode_mask, encode_mask
|
||
|
|
except ModuleNotFoundError as e:
|
||
|
|
if e.name.split(".")[0] != "numpy":
|
||
|
|
raise
|
||
|
|
|
||
|
|
numpy_installed = False
|
||
|
|
else:
|
||
|
|
numpy_installed = True
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.skipif(not numpy_installed, reason="NumPy is not installed")
|
||
|
|
class TestMasks:
|
||
|
|
def _bitmap_from_string(self, string: str) -> np.ndarray:
|
||
|
|
return np.array(
|
||
|
|
[[bool(int(c)) for c in line] for line in textwrap.dedent(string).splitlines()]
|
||
|
|
)
|
||
|
|
|
||
|
|
def test_encode_mask_without_bbox(self):
|
||
|
|
bitmap = self._bitmap_from_string(
|
||
|
|
"""\
|
||
|
|
0000000
|
||
|
|
0001110
|
||
|
|
0011000
|
||
|
|
0000000
|
||
|
|
"""
|
||
|
|
)
|
||
|
|
|
||
|
|
assert encode_mask(bitmap) == [1, 5, 2, 2, 1, 5, 2]
|
||
|
|
|
||
|
|
bitmap = np.zeros((4, 6), dtype=bool)
|
||
|
|
assert encode_mask(bitmap) == [1, 0, 0, 0, 0]
|
||
|
|
|
||
|
|
bitmap = np.ones((4, 6), dtype=bool)
|
||
|
|
assert encode_mask(bitmap) == [0, 4 * 6, 0, 0, 5, 3]
|
||
|
|
|
||
|
|
def test_encode_mask_with_bbox(self):
|
||
|
|
bitmap = self._bitmap_from_string(
|
||
|
|
"""\
|
||
|
|
001110
|
||
|
|
011000
|
||
|
|
"""
|
||
|
|
)
|
||
|
|
bbox = [2.9, 0.9, 4.1, 1.1] # will get rounded to [2, 0, 5, 2]
|
||
|
|
|
||
|
|
# There's slightly different logic for when the cropped mask starts with
|
||
|
|
# 0 and 1, so test both.
|
||
|
|
# This one starts with 1:
|
||
|
|
# 111
|
||
|
|
# 100
|
||
|
|
|
||
|
|
assert encode_mask(bitmap, bbox) == [0, 4, 2, 2, 0, 4, 1]
|
||
|
|
|
||
|
|
bbox = [1, 0, 5, 2]
|
||
|
|
|
||
|
|
# This one starts with 0:
|
||
|
|
# 0111
|
||
|
|
# 1100
|
||
|
|
|
||
|
|
assert encode_mask(bitmap, bbox) == [1, 5, 2, 1, 0, 4, 1]
|
||
|
|
|
||
|
|
bbox = [3, 1, 6, 2] # zeroes only: 000
|
||
|
|
assert encode_mask(bitmap, bbox) == [3, 3, 1, 5, 1]
|
||
|
|
|
||
|
|
bbox = [2, 0, 5, 1] # ones only: 111
|
||
|
|
assert encode_mask(bitmap, bbox) == [0, 3, 2, 0, 4, 0]
|
||
|
|
|
||
|
|
# Edge case: full image
|
||
|
|
bbox = [0, 0, 6, 2]
|
||
|
|
assert encode_mask(bitmap, bbox) == [2, 3, 2, 2, 3, 0, 0, 5, 1]
|
||
|
|
|
||
|
|
def test_encode_mask_invalid_dim(self):
|
||
|
|
with pytest.raises(ValueError, match="bitmap must have 2 dimensions"):
|
||
|
|
encode_mask([True], [0, 0, 1, 1])
|
||
|
|
|
||
|
|
def test_encode_mask_invalid_dtype(self):
|
||
|
|
with pytest.raises(ValueError, match="bitmap must have boolean items"):
|
||
|
|
encode_mask([[1]], [0, 0, 1, 1])
|
||
|
|
|
||
|
|
@pytest.mark.parametrize(
|
||
|
|
"bbox",
|
||
|
|
[
|
||
|
|
[-0.1, 0, 1, 1],
|
||
|
|
[0, -0.1, 1, 1],
|
||
|
|
[0, 0, 1.1, 1],
|
||
|
|
[0, 0, 1, 1.1],
|
||
|
|
[1, 0, 0, 1],
|
||
|
|
[0, 1, 1, 0],
|
||
|
|
[0, 0, 0, 1],
|
||
|
|
[0, 0, 1, 0],
|
||
|
|
],
|
||
|
|
)
|
||
|
|
def test_encode_mask_invalid_bbox(self, bbox):
|
||
|
|
with pytest.raises(ValueError, match="bbox has invalid coordinates"):
|
||
|
|
encode_mask([[True]], bbox)
|
||
|
|
|
||
|
|
def test_decode_mask(self):
|
||
|
|
points = [1, 0, 0, 0, 0]
|
||
|
|
np.testing.assert_array_equal(
|
||
|
|
decode_mask(points, image_width=6, image_height=4), np.zeros((4, 6), dtype=bool)
|
||
|
|
)
|
||
|
|
|
||
|
|
points = [0, 24, 0, 0, 5, 3]
|
||
|
|
np.testing.assert_array_equal(
|
||
|
|
decode_mask(points, image_width=6, image_height=4),
|
||
|
|
np.ones((4, 6), dtype=bool),
|
||
|
|
)
|
||
|
|
|
||
|
|
points = [1, 5, 2, 2, 1, 5, 2]
|
||
|
|
|
||
|
|
bitmap = self._bitmap_from_string(
|
||
|
|
"""\
|
||
|
|
0000000
|
||
|
|
0001110
|
||
|
|
0011000
|
||
|
|
0000000
|
||
|
|
"""
|
||
|
|
)
|
||
|
|
np.testing.assert_array_equal(decode_mask(points, image_width=7, image_height=4), bitmap)
|
||
|
|
|
||
|
|
# Same mask, but with the bbox covering the whole image.
|
||
|
|
points = [10, 3, 3, 2, 10, 0, 0, 6, 3]
|
||
|
|
np.testing.assert_array_equal(decode_mask(points, image_width=7, image_height=4), bitmap)
|
||
|
|
|
||
|
|
@pytest.mark.parametrize("image_width, image_height", [(1, 0), (1, -1), (0, 1), (-1, 1)])
|
||
|
|
def test_decode_mask_invalid_image_dimensions(self, image_width, image_height):
|
||
|
|
with pytest.raises(ValueError, match="invalid image dimensions"):
|
||
|
|
decode_mask([1, 0, 0, 0, 0], image_width=image_width, image_height=image_height)
|
||
|
|
|
||
|
|
@pytest.mark.parametrize("points", [[], [1, 2, 3, 4]])
|
||
|
|
def test_decode_mask_invalid_too_few_elements(self, points):
|
||
|
|
with pytest.raises(ValueError, match="too few elements in encoded mask"):
|
||
|
|
decode_mask(points, image_width=1, image_height=1)
|
||
|
|
|
||
|
|
@pytest.mark.parametrize("points", [[1.1, 2, 3, 4, 5], ["1", 2, 3, 4, 5]])
|
||
|
|
def test_decode_mask_invalid_non_integer(self, points):
|
||
|
|
with pytest.raises(ValueError, match="non-integer value in encoded mask"):
|
||
|
|
decode_mask(points, image_width=1, image_height=1)
|
||
|
|
|
||
|
|
@pytest.mark.parametrize(
|
||
|
|
"points",
|
||
|
|
[
|
||
|
|
[9, -1, 1, 3, 3],
|
||
|
|
[9, 1, -1, 3, 3],
|
||
|
|
[9, 1, 1, 0, 3],
|
||
|
|
[9, 1, 1, 3, 0],
|
||
|
|
[9, 1, 1, 6, 3],
|
||
|
|
[9, 1, 1, 3, 4],
|
||
|
|
],
|
||
|
|
)
|
||
|
|
def test_decode_mask_invalid_bbox(self, points):
|
||
|
|
with pytest.raises(ValueError, match="invalid encoded bounding box"):
|
||
|
|
decode_mask(points, image_width=6, image_height=4)
|
||
|
|
|
||
|
|
def test_decode_mask_invalid_mismatched_bbox(self):
|
||
|
|
with pytest.raises(ValueError, match="encoded bitmap does not match encoded bounding box"):
|
||
|
|
decode_mask([10, 1, 1, 3, 3], image_width=6, image_height=4)
|
||
|
|
|
||
|
|
def _random_subrange(self, rng: np.random.Generator, range_len: int) -> tuple[int, int]:
|
||
|
|
subrange_len = rng.integers(1, range_len + 1)
|
||
|
|
start = rng.integers(0, range_len - subrange_len + 1)
|
||
|
|
return start.item(), (start + subrange_len).item()
|
||
|
|
|
||
|
|
@pytest.mark.parametrize("with_bbox", [False, True])
|
||
|
|
def test_roundtrip(self, with_bbox: bool):
|
||
|
|
rng = np.random.default_rng(seed=list(b"CVAT"))
|
||
|
|
|
||
|
|
for _ in range(100):
|
||
|
|
width, height = rng.integers(1, 11, size=2)
|
||
|
|
bitmap = rng.integers(0, 2, size=(height, width), dtype=bool)
|
||
|
|
|
||
|
|
if with_bbox:
|
||
|
|
x1, x2 = self._random_subrange(rng, width)
|
||
|
|
y1, y2 = self._random_subrange(rng, height)
|
||
|
|
points = encode_mask(bitmap, [x1, y1, x2, y2])
|
||
|
|
expected = np.zeros((height, width), dtype=bool)
|
||
|
|
expected[y1:y2, x1:x2] = bitmap[y1:y2, x1:x2]
|
||
|
|
else:
|
||
|
|
points = encode_mask(bitmap)
|
||
|
|
expected = bitmap
|
||
|
|
|
||
|
|
decoded = decode_mask(points, image_width=width, image_height=height)
|
||
|
|
np.testing.assert_array_equal(decoded, expected)
|