330 lines
12 KiB
Python
330 lines
12 KiB
Python
import unittest
|
|
|
|
try:
|
|
from .test_mask_draw_helpers import load_plugin_package, torch
|
|
except ImportError:
|
|
from test_mask_draw_helpers import load_plugin_package, torch
|
|
|
|
|
|
HAS_CONCAT_RUNTIME_DEPS = torch is not None
|
|
CONCAT_RUNTIME_DEPS_MESSAGE = "concat node tests require torch"
|
|
|
|
|
|
@unittest.skipUnless(HAS_CONCAT_RUNTIME_DEPS, CONCAT_RUNTIME_DEPS_MESSAGE)
|
|
class TestConcatByTargetNode(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
plugin = load_plugin_package()
|
|
cls.node = plugin.NODE_CLASS_MAPPINGS["LLSConcatByTarget"]()
|
|
|
|
def make_image(self, width, height, value, batch=1):
|
|
return torch.full((batch, height, width, 3), float(value), dtype=torch.float32)
|
|
|
|
def make_mask(self, width, height, value, batch=1):
|
|
return torch.full((batch, height, width), float(value), dtype=torch.float32)
|
|
|
|
def assert_image_output(self, image, mask, width, height, shape):
|
|
self.assertEqual((width, height), (shape[2], shape[1]))
|
|
self.assertEqual(tuple(image.shape), shape)
|
|
self.assertEqual(tuple(mask.shape), (shape[0], shape[1], shape[2]))
|
|
self.assertGreaterEqual(float(image.min().item()), 0.0)
|
|
self.assertLessEqual(float(image.max().item()), 1.0)
|
|
self.assertTrue(torch.allclose(mask, torch.zeros_like(mask)))
|
|
|
|
def test_legacy_port_names_are_accepted(self):
|
|
image_a = self.make_image(width=2, height=2, value=0.2)
|
|
image_b = self.make_image(width=1, height=2, value=0.8)
|
|
|
|
image, mask, width, height = self.node.concat(
|
|
data_type="IMAGE",
|
|
target="A",
|
|
position="right",
|
|
match_target_size=True,
|
|
resize_mode="none",
|
|
align="center",
|
|
gap=0,
|
|
background_color="#000000",
|
|
background_value=0.0,
|
|
multiple_of=0,
|
|
allow_batch_broadcast=True,
|
|
**{"image/mask_A": image_a, "image/mask_B": image_b},
|
|
)
|
|
|
|
self.assert_image_output(image, mask, width, height, (1, 2, 3, 3))
|
|
self.assertTrue(torch.allclose(image[:, :, 0:2, :], torch.full((1, 2, 2, 3), 0.2)))
|
|
self.assertTrue(torch.allclose(image[:, :, 2:3, :], torch.full((1, 2, 1, 3), 0.8)))
|
|
|
|
def test_image_mode_target_b_right_places_a_to_the_right_of_b(self):
|
|
image_a = self.make_image(width=2, height=2, value=0.2)
|
|
image_b = self.make_image(width=3, height=2, value=0.8)
|
|
|
|
image, mask, width, height = self.node.concat(
|
|
data_type="IMAGE",
|
|
target="B",
|
|
position="right",
|
|
match_target_size=True,
|
|
resize_mode="none",
|
|
align="center",
|
|
gap=1,
|
|
background_color="#000000",
|
|
background_value=0.0,
|
|
multiple_of=0,
|
|
allow_batch_broadcast=True,
|
|
a=image_a,
|
|
b=image_b,
|
|
)
|
|
|
|
self.assert_image_output(image, mask, width, height, (1, 2, 6, 3))
|
|
self.assertTrue(torch.allclose(image[:, :, 0:3, :], torch.full((1, 2, 3, 3), 0.8)))
|
|
self.assertTrue(torch.allclose(image[:, :, 3:4, :], torch.zeros((1, 2, 1, 3))))
|
|
self.assertTrue(torch.allclose(image[:, :, 4:6, :], torch.full((1, 2, 2, 3), 0.2)))
|
|
|
|
def test_image_mode_keep_proportion_gap_and_multiple_of(self):
|
|
image_a = self.make_image(width=4, height=4, value=0.25)
|
|
image_b = self.make_image(width=2, height=2, value=0.75)
|
|
|
|
image, mask, width, height = self.node.concat(
|
|
data_type="IMAGE",
|
|
target="A",
|
|
position="right",
|
|
match_target_size=True,
|
|
resize_mode="keep_proportion",
|
|
align="center",
|
|
gap=1,
|
|
background_color="#808080",
|
|
background_value=0.0,
|
|
multiple_of=4,
|
|
allow_batch_broadcast=True,
|
|
a=image_a,
|
|
b=image_b,
|
|
)
|
|
|
|
self.assert_image_output(image, mask, width, height, (1, 4, 12, 3))
|
|
self.assertTrue(torch.allclose(image[:, :, 0:4, :], torch.full((1, 4, 4, 3), 0.25)))
|
|
self.assertTrue(torch.allclose(image[:, :, 4:5, :], torch.full((1, 4, 1, 3), 128.0 / 255.0)))
|
|
self.assertTrue(torch.allclose(image[:, :, 5:9, :], torch.full((1, 4, 4, 3), 0.75)))
|
|
self.assertTrue(torch.allclose(image[:, :, 9:12, :], torch.full((1, 4, 3, 3), 128.0 / 255.0)))
|
|
|
|
def test_mask_mode_outputs_mask_and_preview_image(self):
|
|
mask_a = self.make_mask(width=2, height=2, value=1.0)
|
|
mask_b = self.make_mask(width=2, height=2, value=0.0)
|
|
|
|
image, mask, width, height = self.node.concat(
|
|
data_type="MASK",
|
|
target="A",
|
|
position="bottom",
|
|
match_target_size=True,
|
|
resize_mode="none",
|
|
align="center",
|
|
gap=1,
|
|
background_color="#000000",
|
|
background_value=0.25,
|
|
multiple_of=0,
|
|
allow_batch_broadcast=True,
|
|
a=mask_a,
|
|
b=mask_b,
|
|
)
|
|
|
|
self.assertEqual((width, height), (2, 5))
|
|
self.assertEqual(tuple(mask.shape), (1, 5, 2))
|
|
self.assertEqual(tuple(image.shape), (1, 5, 2, 3))
|
|
self.assertTrue(torch.allclose(mask[:, 0:2, :], torch.full((1, 2, 2), 1.0)))
|
|
self.assertTrue(torch.allclose(mask[:, 2:3, :], torch.full((1, 1, 2), 0.25)))
|
|
self.assertTrue(torch.allclose(mask[:, 3:5, :], torch.zeros((1, 2, 2))))
|
|
self.assertTrue(torch.allclose(image[..., 0], mask))
|
|
self.assertTrue(torch.allclose(image[..., 1], mask))
|
|
self.assertTrue(torch.allclose(image[..., 2], mask))
|
|
|
|
def test_batch_broadcast_works_when_enabled(self):
|
|
image_a = self.make_image(width=2, height=2, value=0.1, batch=1)
|
|
image_b = self.make_image(width=2, height=2, value=0.9, batch=2)
|
|
|
|
image, mask, width, height = self.node.concat(
|
|
data_type="IMAGE",
|
|
target="A",
|
|
position="left",
|
|
match_target_size=True,
|
|
resize_mode="none",
|
|
align="center",
|
|
gap=0,
|
|
background_color="#000000",
|
|
background_value=0.0,
|
|
multiple_of=0,
|
|
allow_batch_broadcast=True,
|
|
a=image_a,
|
|
b=image_b,
|
|
)
|
|
|
|
self.assert_image_output(image, mask, width, height, (2, 2, 4, 3))
|
|
self.assertTrue(torch.allclose(image[:, :, 0:2, :], torch.full((2, 2, 2, 3), 0.9)))
|
|
self.assertTrue(torch.allclose(image[:, :, 2:4, :], torch.full((2, 2, 2, 3), 0.1)))
|
|
|
|
def test_gapless_concat_keeps_output_compact_when_sizes_differ(self):
|
|
image_a = self.make_image(width=2, height=2, value=0.1)
|
|
image_b = self.make_image(width=3, height=4, value=0.9)
|
|
|
|
image, mask, width, height = self.node.concat(
|
|
data_type="IMAGE",
|
|
target="A",
|
|
position="right",
|
|
match_target_size=False,
|
|
resize_mode="none",
|
|
align="center",
|
|
gap=0,
|
|
background_color="#000000",
|
|
background_value=0.0,
|
|
multiple_of=0,
|
|
allow_batch_broadcast=True,
|
|
a=image_a,
|
|
b=image_b,
|
|
)
|
|
|
|
self.assert_image_output(image, mask, width, height, (1, 2, 5, 3))
|
|
self.assertTrue(torch.allclose(image[:, :, 0:2, :], torch.full((1, 2, 2, 3), 0.1)))
|
|
self.assertTrue(torch.allclose(image[:, :, 2:5, :], torch.full((1, 2, 3, 3), 0.9)))
|
|
|
|
def test_mask_inputs_are_auto_detected_even_if_data_type_is_image(self):
|
|
mask_a = self.make_mask(width=2, height=2, value=1.0)
|
|
mask_b = self.make_mask(width=1, height=2, value=0.5)
|
|
|
|
image, mask, width, height = self.node.concat(
|
|
data_type="IMAGE",
|
|
target="A",
|
|
position="right",
|
|
match_target_size=True,
|
|
resize_mode="none",
|
|
align="center",
|
|
gap=0,
|
|
background_color="#000000",
|
|
background_value=0.0,
|
|
multiple_of=0,
|
|
allow_batch_broadcast=True,
|
|
a=mask_a,
|
|
b=mask_b,
|
|
)
|
|
|
|
self.assertEqual((width, height), (3, 2))
|
|
self.assertEqual(tuple(mask.shape), (1, 2, 3))
|
|
self.assertEqual(tuple(image.shape), (1, 2, 3, 3))
|
|
self.assertTrue(torch.allclose(mask[:, :, 0:2], torch.ones((1, 2, 2))))
|
|
self.assertTrue(torch.allclose(mask[:, :, 2:3], torch.full((1, 2, 1), 0.5)))
|
|
|
|
def test_image_inputs_are_auto_detected_even_if_data_type_is_mask(self):
|
|
image_a = self.make_image(width=2, height=2, value=0.2)
|
|
image_b = self.make_image(width=1, height=2, value=0.8)
|
|
|
|
image, mask, width, height = self.node.concat(
|
|
data_type="MASK",
|
|
target="A",
|
|
position="right",
|
|
match_target_size=True,
|
|
resize_mode="none",
|
|
align="center",
|
|
gap=0,
|
|
background_color="#000000",
|
|
background_value=0.0,
|
|
multiple_of=0,
|
|
allow_batch_broadcast=True,
|
|
a=image_a,
|
|
b=image_b,
|
|
)
|
|
|
|
self.assert_image_output(image, mask, width, height, (1, 2, 3, 3))
|
|
self.assertTrue(torch.allclose(image[:, :, 0:2, :], torch.full((1, 2, 2, 3), 0.2)))
|
|
self.assertTrue(torch.allclose(image[:, :, 2:3, :], torch.full((1, 2, 1, 3), 0.8)))
|
|
|
|
def test_mixed_image_and_mask_inputs_raise_clear_error(self):
|
|
with self.assertRaisesRegex(RuntimeError, "must both resolve to IMAGE or both resolve to MASK"):
|
|
self.node.concat(
|
|
data_type="IMAGE",
|
|
target="A",
|
|
position="right",
|
|
match_target_size=True,
|
|
resize_mode="none",
|
|
align="center",
|
|
gap=0,
|
|
background_color="#000000",
|
|
background_value=0.0,
|
|
multiple_of=0,
|
|
allow_batch_broadcast=True,
|
|
a=self.make_image(2, 2, 0.2),
|
|
b=self.make_mask(2, 2, 1.0),
|
|
)
|
|
|
|
def test_missing_required_inputs_for_current_mode_raise_clear_errors(self):
|
|
with self.assertRaisesRegex(RuntimeError, "image/mask_A and image/mask_B"):
|
|
self.node.concat(
|
|
data_type="IMAGE",
|
|
target="A",
|
|
position="right",
|
|
match_target_size=True,
|
|
resize_mode="none",
|
|
align="center",
|
|
gap=0,
|
|
background_color="#000000",
|
|
background_value=0.0,
|
|
multiple_of=0,
|
|
allow_batch_broadcast=True,
|
|
a=None,
|
|
b=None,
|
|
)
|
|
|
|
with self.assertRaisesRegex(RuntimeError, "image/mask_A and image/mask_B"):
|
|
self.node.concat(
|
|
data_type="MASK",
|
|
target="A",
|
|
position="right",
|
|
match_target_size=True,
|
|
resize_mode="none",
|
|
align="center",
|
|
gap=0,
|
|
background_color="#000000",
|
|
background_value=0.0,
|
|
multiple_of=0,
|
|
allow_batch_broadcast=True,
|
|
a=None,
|
|
b=None,
|
|
)
|
|
|
|
def test_invalid_background_color_and_batch_mismatch_raise_clear_errors(self):
|
|
image_a = self.make_image(width=2, height=2, value=0.1, batch=2)
|
|
image_b = self.make_image(width=2, height=2, value=0.9, batch=3)
|
|
|
|
with self.assertRaisesRegex(RuntimeError, "background_color"):
|
|
self.node.concat(
|
|
data_type="IMAGE",
|
|
target="A",
|
|
position="right",
|
|
match_target_size=True,
|
|
resize_mode="none",
|
|
align="center",
|
|
gap=0,
|
|
background_color="bad-color",
|
|
background_value=0.0,
|
|
multiple_of=0,
|
|
allow_batch_broadcast=True,
|
|
a=self.make_image(2, 2, 0.1),
|
|
b=self.make_image(2, 2, 0.9),
|
|
)
|
|
|
|
with self.assertRaisesRegex(RuntimeError, "Batch size mismatch"):
|
|
self.node.concat(
|
|
data_type="IMAGE",
|
|
target="A",
|
|
position="right",
|
|
match_target_size=True,
|
|
resize_mode="none",
|
|
align="center",
|
|
gap=0,
|
|
background_color="#000000",
|
|
background_value=0.0,
|
|
multiple_of=0,
|
|
allow_batch_broadcast=False,
|
|
a=image_a,
|
|
b=image_b,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|