修改图像拼接节点名字
This commit is contained in:
@@ -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",
|
||||
|
||||
@@ -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)
|
||||
|
||||
+22
-13
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user