183 lines
6.7 KiB
Python
183 lines
6.7 KiB
Python
import json
|
|
import unittest
|
|
|
|
try:
|
|
from .test_mask_draw_helpers import (
|
|
HAS_MASK_DRAW_RUNTIME_DEPS,
|
|
MASK_DRAW_RUNTIME_DEPS_MESSAGE,
|
|
load_mask_draw_utils,
|
|
make_image,
|
|
make_mask,
|
|
make_mask_state_json,
|
|
torch,
|
|
)
|
|
except ImportError:
|
|
from test_mask_draw_helpers import (
|
|
HAS_MASK_DRAW_RUNTIME_DEPS,
|
|
MASK_DRAW_RUNTIME_DEPS_MESSAGE,
|
|
load_mask_draw_utils,
|
|
make_image,
|
|
make_mask,
|
|
make_mask_state_json,
|
|
torch,
|
|
)
|
|
|
|
|
|
class TestMaskDrawUtils(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.utils = load_mask_draw_utils()
|
|
|
|
def test_parse_mask_state_defaults_invalid_json(self):
|
|
state = self.utils.parse_mask_state("{invalid")
|
|
|
|
self.assertFalse(state["touched"])
|
|
self.assertEqual(state["mask_png_base64"], "")
|
|
self.assertEqual(state["version"], 1)
|
|
|
|
def test_parse_mask_state_defaults_invalid_editor_and_dirty_touched_values(self):
|
|
for payload in (
|
|
{"editor": [1], "touched": "false"},
|
|
{"editor": "bad", "touched": "true"},
|
|
):
|
|
with self.subTest(payload=payload):
|
|
state = self.utils.parse_mask_state(json.dumps(payload))
|
|
|
|
self.assertEqual(state["editor"], {})
|
|
self.assertFalse(state["touched"])
|
|
|
|
@unittest.skipUnless(HAS_MASK_DRAW_RUNTIME_DEPS, MASK_DRAW_RUNTIME_DEPS_MESSAGE)
|
|
def test_resolve_output_mask_prefers_saved_mask_when_touched(self):
|
|
image = make_image(width=6, height=4)
|
|
input_mask = make_mask(width=6, height=4, value=0.0)
|
|
state_json = make_mask_state_json(width=6, height=4, value=1.0, touched=True)
|
|
|
|
mask = self.utils.resolve_output_mask(
|
|
image=image,
|
|
input_mask=input_mask,
|
|
mask_state_json=state_json,
|
|
invert_mask=False,
|
|
)
|
|
|
|
self.assertEqual(tuple(mask.shape), (1, 4, 6))
|
|
self.assertTrue(torch.all(mask == 1.0))
|
|
|
|
@unittest.skipUnless(HAS_MASK_DRAW_RUNTIME_DEPS, MASK_DRAW_RUNTIME_DEPS_MESSAGE)
|
|
def test_resolve_output_mask_repeats_saved_mask_to_image_batch(self):
|
|
image = torch.full((3, 4, 6, 3), 0.25, dtype=torch.float32)
|
|
state_json = make_mask_state_json(width=6, height=4, value=1.0, touched=True)
|
|
|
|
mask = self.utils.resolve_output_mask(
|
|
image=image,
|
|
input_mask=None,
|
|
mask_state_json=state_json,
|
|
invert_mask=False,
|
|
)
|
|
|
|
self.assertEqual(tuple(mask.shape), (3, 4, 6))
|
|
self.assertTrue(torch.all(mask == 1.0))
|
|
|
|
@unittest.skipUnless(HAS_MASK_DRAW_RUNTIME_DEPS, MASK_DRAW_RUNTIME_DEPS_MESSAGE)
|
|
def test_resolve_output_mask_uses_input_mask_when_untouched(self):
|
|
image = make_image(width=6, height=4)
|
|
input_mask = make_mask(width=3, height=2, value=1.0)
|
|
|
|
mask = self.utils.resolve_output_mask(
|
|
image=image,
|
|
input_mask=input_mask,
|
|
mask_state_json="{}",
|
|
invert_mask=False,
|
|
)
|
|
|
|
self.assertEqual(tuple(mask.shape), (1, 4, 6))
|
|
self.assertTrue(torch.all(mask == 1.0))
|
|
|
|
@unittest.skipUnless(HAS_MASK_DRAW_RUNTIME_DEPS, MASK_DRAW_RUNTIME_DEPS_MESSAGE)
|
|
def test_resolve_output_mask_repeats_single_input_mask_to_image_batch(self):
|
|
image = torch.full((3, 4, 6, 3), 0.25, dtype=torch.float32)
|
|
input_mask = make_mask(width=6, height=4, value=1.0)
|
|
|
|
mask = self.utils.resolve_output_mask(
|
|
image=image,
|
|
input_mask=input_mask,
|
|
mask_state_json="{}",
|
|
invert_mask=False,
|
|
)
|
|
|
|
self.assertEqual(tuple(mask.shape), (3, 4, 6))
|
|
self.assertTrue(torch.all(mask == 1.0))
|
|
|
|
@unittest.skipUnless(HAS_MASK_DRAW_RUNTIME_DEPS, MASK_DRAW_RUNTIME_DEPS_MESSAGE)
|
|
def test_resolve_output_mask_rejects_input_mask_batch_mismatch(self):
|
|
image = torch.full((3, 4, 6, 3), 0.25, dtype=torch.float32)
|
|
input_mask = torch.full((2, 4, 6), 1.0, dtype=torch.float32)
|
|
|
|
with self.assertRaisesRegex(RuntimeError, "mask batch size must match image batch size"):
|
|
self.utils.resolve_output_mask(
|
|
image=image,
|
|
input_mask=input_mask,
|
|
mask_state_json="{}",
|
|
invert_mask=False,
|
|
)
|
|
|
|
@unittest.skipUnless(HAS_MASK_DRAW_RUNTIME_DEPS, MASK_DRAW_RUNTIME_DEPS_MESSAGE)
|
|
def test_resolve_output_mask_returns_black_after_clear_state(self):
|
|
image = make_image(width=6, height=4)
|
|
input_mask = make_mask(width=6, height=4, value=1.0)
|
|
state_json = make_mask_state_json(width=6, height=4, value=0.0, touched=True)
|
|
|
|
mask = self.utils.resolve_output_mask(
|
|
image=image,
|
|
input_mask=input_mask,
|
|
mask_state_json=state_json,
|
|
invert_mask=False,
|
|
)
|
|
|
|
self.assertTrue(torch.all(mask == 0.0))
|
|
|
|
@unittest.skipUnless(HAS_MASK_DRAW_RUNTIME_DEPS, MASK_DRAW_RUNTIME_DEPS_MESSAGE)
|
|
def test_resolve_output_mask_falls_back_to_input_mask_when_saved_mask_is_corrupt(self):
|
|
image = make_image(width=6, height=4)
|
|
input_mask = make_mask(width=6, height=4, value=1.0)
|
|
state_json = '{"version": 1, "mask_png_base64": "not-base64", "touched": true, "editor": {}}'
|
|
|
|
mask = self.utils.resolve_output_mask(
|
|
image=image,
|
|
input_mask=input_mask,
|
|
mask_state_json=state_json,
|
|
invert_mask=False,
|
|
)
|
|
|
|
self.assertEqual(tuple(mask.shape), (1, 4, 6))
|
|
self.assertTrue(torch.all(mask == 1.0))
|
|
|
|
@unittest.skipUnless(HAS_MASK_DRAW_RUNTIME_DEPS, MASK_DRAW_RUNTIME_DEPS_MESSAGE)
|
|
def test_resolve_output_mask_applies_invert(self):
|
|
image = make_image(width=6, height=4)
|
|
state_json = make_mask_state_json(width=6, height=4, value=1.0, touched=True)
|
|
|
|
mask = self.utils.resolve_output_mask(
|
|
image=image,
|
|
input_mask=None,
|
|
mask_state_json=state_json,
|
|
invert_mask=True,
|
|
)
|
|
|
|
self.assertTrue(torch.all(mask == 0.0))
|
|
|
|
@unittest.skipUnless(HAS_MASK_DRAW_RUNTIME_DEPS, MASK_DRAW_RUNTIME_DEPS_MESSAGE)
|
|
def test_build_preview_image_preserves_shape_and_adds_red_overlay(self):
|
|
image = make_image(width=6, height=4, color=0.2)
|
|
mask = make_mask(width=6, height=4, value=1.0)
|
|
|
|
preview = self.utils.build_preview_image(image=image, mask=mask, overlay_alpha=0.4)
|
|
|
|
self.assertEqual(tuple(preview.shape), (1, 4, 6, 3))
|
|
self.assertGreater(float(preview[..., 0].mean()), float(image[..., 0].mean()))
|
|
self.assertLess(float(preview[..., 1].mean()), float(image[..., 1].mean()))
|
|
self.assertLess(float(preview[..., 2].mean()), float(image[..., 2].mean()))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|