diff --git a/tests/test_concat_by_target_node.py b/tests/test_concat_by_target_node.py index 5ec619f..bead86a 100644 --- a/tests/test_concat_by_target_node.py +++ b/tests/test_concat_by_target_node.py @@ -31,6 +31,29 @@ class TestConcatByTargetNode(unittest.TestCase): self.assertLessEqual(float(image.max().item()), 1.0) self.assertTrue(torch.allclose(mask, torch.zeros_like(mask))) + def test_new_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) @@ -229,7 +252,7 @@ class TestConcatByTargetNode(unittest.TestCase): ) def test_missing_required_inputs_for_current_mode_raise_clear_errors(self): - with self.assertRaisesRegex(RuntimeError, "inputs a and b"): + with self.assertRaisesRegex(RuntimeError, "image/mask_A and image/mask_B"): self.node.concat( data_type="IMAGE", target="A", @@ -246,7 +269,7 @@ class TestConcatByTargetNode(unittest.TestCase): b=None, ) - with self.assertRaisesRegex(RuntimeError, "inputs a and b"): + with self.assertRaisesRegex(RuntimeError, "image/mask_A and image/mask_B"): self.node.concat( data_type="MASK", target="A", diff --git a/tests/test_concat_by_target_registration.py b/tests/test_concat_by_target_registration.py index 86e1db6..e1d6965 100644 --- a/tests/test_concat_by_target_registration.py +++ b/tests/test_concat_by_target_registration.py @@ -41,8 +41,10 @@ class TestConcatByTargetRegistration(unittest.TestCase): self.assertEqual(required["background_value"][1]["default"], 0.0) self.assertEqual(required["multiple_of"][0], "INT") self.assertEqual(required["allow_batch_broadcast"][0], "BOOLEAN") - self.assertEqual(optional["a"][0], "IMAGE,MASK") - self.assertEqual(optional["b"][0], "IMAGE,MASK") + self.assertEqual(optional["image/mask_A"][0], "IMAGE,MASK") + self.assertEqual(optional["image/mask_B"][0], "IMAGE,MASK") + self.assertNotIn("a", optional) + self.assertNotIn("b", optional) self.assertNotIn("image_a", optional) self.assertNotIn("image_b", optional) self.assertNotIn("mask_a", optional) diff --git a/utils/concat_by_target.py b/utils/concat_by_target.py index 13dc0f6..5e49b81 100644 --- a/utils/concat_by_target.py +++ b/utils/concat_by_target.py @@ -17,6 +17,8 @@ POSITION_CHOICES = ["top", "bottom", "left", "right"] RESIZE_MODE_CHOICES = ["keep_proportion", "stretch", "none"] ALIGN_CHOICES = ["start", "center", "end"] INPUT_DATA_CHOICES = "IMAGE,MASK" +PORT_A_NAME = "image/mask_A" +PORT_B_NAME = "image/mask_B" def _require_torch(): @@ -57,8 +59,8 @@ def resolve_concat_data_type(a, b, requested_data_type: str): if len(inferred) > 1: raise RuntimeError( - f"[LLS] Inputs a and b must both resolve to IMAGE or both resolve to MASK. " - f"Detected a={inferred_a or 'unknown'}, b={inferred_b or 'unknown'}." + f"[LLS] Inputs {PORT_A_NAME} and {PORT_B_NAME} must both resolve to IMAGE or both resolve to MASK. " + f"Detected {PORT_A_NAME}={inferred_a or 'unknown'}, {PORT_B_NAME}={inferred_b or 'unknown'}." ) if len(inferred) == 1: return inferred.pop() @@ -450,8 +452,8 @@ class LLSConcatByTarget: "allow_batch_broadcast": ("BOOLEAN", {"default": True}), }, "optional": { - "a": (INPUT_DATA_CHOICES,), - "b": (INPUT_DATA_CHOICES,), + PORT_A_NAME: (INPUT_DATA_CHOICES,), + PORT_B_NAME: (INPUT_DATA_CHOICES,), }, } @@ -470,18 +472,23 @@ class LLSConcatByTarget: allow_batch_broadcast, a=None, b=None, + **kwargs, ): _require_torch() data_type = str(data_type or "IMAGE") background_value = max(0.0, min(1.0, float(background_value))) - resolved_data_type = resolve_concat_data_type(a, b, data_type) + input_a = kwargs[PORT_A_NAME] if PORT_A_NAME in kwargs else a + input_b = kwargs[PORT_B_NAME] if PORT_B_NAME in kwargs else b + resolved_data_type = resolve_concat_data_type(input_a, input_b, data_type) if resolved_data_type == "IMAGE": - if a is None or b is None: - raise RuntimeError("[LLS] IMAGE mode requires both inputs a and b to be connected.") + if input_a is None or input_b is None: + raise RuntimeError( + f"[LLS] IMAGE mode requires both inputs {PORT_A_NAME} and {PORT_B_NAME} to be connected." + ) color = parse_hex_color(background_color) - tensor_a = ensure_image_tensor(a) - tensor_b = ensure_image_tensor(b) + tensor_a = ensure_image_tensor(input_a) + tensor_b = ensure_image_tensor(input_b) output_image, width, height = concat_by_target( data_type="IMAGE", tensor_a=tensor_a, @@ -506,10 +513,12 @@ class LLSConcatByTarget: return output_image, output_mask, int(width), int(height) if resolved_data_type == "MASK": - if a is None or b is None: - raise RuntimeError("[LLS] MASK mode requires both inputs a and b to be connected.") - tensor_a = ensure_mask_tensor(a) - tensor_b = ensure_mask_tensor(b) + if input_a is None or input_b is None: + raise RuntimeError( + f"[LLS] MASK mode requires both inputs {PORT_A_NAME} and {PORT_B_NAME} to be connected." + ) + tensor_a = ensure_mask_tensor(input_a) + tensor_b = ensure_mask_tensor(input_b) output_mask, width, height = concat_by_target( data_type="MASK", tensor_a=tensor_a,