diff --git a/README.md b/README.md index 90b8f56..a18055a 100644 --- a/README.md +++ b/README.md @@ -81,6 +81,7 @@ node/ `LLS-node` 现在提供一套面向本地工作流的修复链路，由以下节点组成： - `LLS Simple Mask Create` +- `LLS Simple Mask Preview` - `LLS Simple Mask Draw` - `LLS Simple Repair Prepare` - `LLS Simple KSampler` @@ -90,28 +91,30 @@ node/ ### `LLS Simple Mask Create` -`LLS Simple Mask Create` 是一个面向局部重绘的规则几何遮罩节点，用来根据输入图像尺寸快速创建 `rectangle`、`square`、`circle`、`ellipse` 四种基础 `mask`，并输出面积信息与叠加预览图。 +`LLS Simple Mask Create` 是一个基础几何遮罩生成节点。它不依赖原图内容，只根据输出尺寸和几何参数直接创建 `rectangle`、`square`、`circle`、`ellipse` 四种基础 `mask`。 -它适合放在手工绘制前或修复准备前： +它适合放在修复准备前或手工绘制前： -- `Load Image -> LLS Simple Mask Create -> Preview Image` -- `Load Image -> LLS Simple Mask Create -> LLS Simple Repair Prepare` -- `Load Image -> LLS Simple Mask Create -> LLS Simple Mask Draw` -- `Load Image -> LLS Simple Mask Create -> LLS Simple Mask Draw -> LLS Simple Repair Prepare` +- `LLS Simple Mask Create.mask_image -> Preview Image` +- `Load Image.image + LLS Simple Mask Create.mask -> LLS Simple Mask Preview -> Preview Image` +- `Load Image.image + LLS Simple Mask Create.mask -> LLS Simple Repair Prepare` +- `LLS Simple Mask Create.mask -> LLS Simple Mask Draw.input_mask` **推荐连接方式：** -- `Load Image.image -> LLS Simple Mask Create.image` -- `LLS Simple Mask Create.image -> LLS Simple Repair Prepare.image` +- `LLS Simple Mask Create.mask_image -> Preview Image` +- `Load Image.image + LLS Simple Mask Create.mask -> LLS Simple Mask Preview` +- `Load Image.image -> LLS Simple Repair Prepare.image` - `LLS Simple Mask Create.mask -> LLS Simple Repair Prepare.mask` - `LLS Simple Mask Create.mask -> LLS Simple Mask Draw.input_mask` -- `LLS Simple Mask Create.image -> LLS Simple Mask Draw.image` +- `Load Image.image -> LLS Simple Mask Draw.image` **输入：** | 名称 | 类型 | 说明 | |------|------|------| -| `image` | `IMAGE` | 必填，原图输入，用于确定几何 mask 尺寸 | +| `image_width` | `INT` | 必填，输出遮罩宽度 | +| `image_height` | `INT` | 必填，输出遮罩高度 | | `input_mask` | `MASK` | 可选，已有遮罩；可与当前几何 mask 做 `replace` / `union` / `subtract` / `intersect` | **核心参数：** @@ -127,16 +130,13 @@ node/ | `blur` | 整体边缘模糊 | | `invert_mask` | 反转最终 mask | | `combine_mode` | `replace` / `union` / `subtract` / `intersect` | -| `overlay_alpha` | 预览叠加透明度 | -| `overlay_color` | `red` / `green` / `blue` | **输出：** | 名称 | 类型 | 说明 | |------|------|------| -| `image` | `IMAGE` | 原图透传 | | `mask` | `MASK` | 最终遮罩，白色 / `1` 表示要重绘，黑色 / `0` 表示不重绘 | -| `preview_image` | `IMAGE` | 原图与半透明彩色 mask 叠加后的预览图 | +| `mask_image` | `IMAGE` | 将 `mask` 直接可视化为黑白图像，方便单独预览 | | `area_info` | `LLS_MASK_INFO` | 包含面积、占比、bbox、坐标模式等信息的 dict | `area_info` 主要包含： @@ -165,12 +165,12 @@ node/ - `coordinate_mode = percent` - `center_x = 0.5` - `center_y = 0.5` +- `image_width = 1024` +- `image_height = 1024` - `width = 0.3` - `height = 0.3` - `radius = 0.15` - `combine_mode = replace` -- `overlay_alpha = 0.4` -- `overlay_color = red` 这会在图像中心创建一个大约占宽高 `30%` 的矩形 mask。 @@ -188,6 +188,30 @@ node/ - 先创建基础几何区域，再交给 `LLS Simple Mask Draw` 手动修边 - 为 `LLS Simple Repair Prepare` 提供稳定的初始局部重绘 mask +### `LLS Simple Mask Preview` + +`LLS Simple Mask Preview` 只负责把 `image + mask` 做半透明叠加预览，不参与采样，也不修改原图。 + +**推荐工作流：** + +- `Load Image.image + LLS Simple Mask Create.mask -> LLS Simple Mask Preview -> Preview Image` +- `Load Image.image + LLS Simple Mask Draw.mask -> LLS Simple Mask Preview -> Preview Image` + +**输入：** + +| 名称 | 类型 | 说明 | +|------|------|------| +| `image` | `IMAGE` | 必填，原图输入 | +| `mask` | `MASK` | 必填，待预览的遮罩 | +| `overlay_alpha` | `FLOAT` | 叠加透明度 | +| `overlay_color` | `red / green / blue` | 叠加颜色 | + +**输出：** + +| 名称 | 类型 | 说明 | +|------|------|------| +| `preview_image` | `IMAGE` | 原图与半透明彩色 mask 叠加后的预览图 | + ### `LLS Simple Mask Draw` `LLS Simple Mask Draw` 是一个面向局部重绘的交互式遮罩输入节点，用来直接在输入图像上手动画 `mask`，并把结果输出给后续修复链路。 diff --git a/mask/mask_create.py b/mask/mask_create.py index 9ead438..a130835 100644 --- a/mask/mask_create.py +++ b/mask/mask_create.py @@ -4,29 +4,39 @@ from .mask_utils import ( COMBINE_MODE_CHOICES, COORDINATE_MODE_CHOICES, MASK_INFO_TYPE, - OVERLAY_COLOR_CHOICES, SHAPE_TYPE_CHOICES, apply_invert, apply_softening, build_area_info, - build_preview_image, combine_masks, + create_reference_image, create_shape_mask, + mask_to_image, ) +def _resolve_batch_for_mask_creation(input_mask) -> int: + if input_mask is None: + return 1 + shape = getattr(input_mask, "shape", None) + if not isinstance(shape, (list, tuple)) or len(shape) != 3: + raise RuntimeError("[LLS] mask must have shape [batch, height, width].") + return max(1, int(shape[0])) + + class LLSSimpleMaskCreate: CATEGORY = "LLS/Mask" FUNCTION = "create_mask" - RETURN_TYPES = ("IMAGE", "MASK", "IMAGE", MASK_INFO_TYPE) - RETURN_NAMES = ("image", "mask", "preview_image", "area_info") - DESCRIPTION = "Create a geometric repair mask, preview overlay, and area metadata from an input image." + RETURN_TYPES = ("MASK", "IMAGE", MASK_INFO_TYPE) + RETURN_NAMES = ("mask", "mask_image", "area_info") + DESCRIPTION = "Create a geometric mask image from width, height, and shape parameters." @classmethod def INPUT_TYPES(cls): return { "required": { - "image": ("IMAGE",), + "image_width": ("INT", {"default": 1024, "min": 1, "max": 8192, "step": 1}), + "image_height": ("INT", {"default": 1024, "min": 1, "max": 8192, "step": 1}), "shape_type": (SHAPE_TYPE_CHOICES, {"default": "rectangle"}), "coordinate_mode": (COORDINATE_MODE_CHOICES, {"default": "percent"}), "center_x": ("FLOAT", {"default": 0.5, "step": 0.01}), @@ -38,8 +48,6 @@ class LLSSimpleMaskCreate: "blur": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 128.0, "step": 0.5}), "invert_mask": ("BOOLEAN", {"default": False}), "combine_mode": (COMBINE_MODE_CHOICES, {"default": "replace"}), - "overlay_alpha": ("FLOAT", {"default": 0.4, "min": 0.0, "max": 1.0, "step": 0.01}), - "overlay_color": (OVERLAY_COLOR_CHOICES, {"default": "red"}), }, "optional": { "input_mask": ("MASK",), @@ -48,7 +56,8 @@ class LLSSimpleMaskCreate: def create_mask( self, - image, + image_width, + image_height, shape_type, coordinate_mode, center_x, @@ -60,15 +69,19 @@ class LLSSimpleMaskCreate: blur, invert_mask, combine_mode, - overlay_alpha, - overlay_color, input_mask=None, ): - if image is None: - raise RuntimeError("[LLS] Missing image for LLS Simple Mask Create.") + batch = _resolve_batch_for_mask_creation(input_mask) + reference_image = create_reference_image( + image_width=int(image_width), + image_height=int(image_height), + batch=batch, + dtype=getattr(input_mask, "dtype", None), + device=getattr(input_mask, "device", None), + ) created_mask, geometry_info = create_shape_mask( - image, + reference_image, shape_type=shape_type, coordinate_mode=coordinate_mode, center_x=center_x, @@ -77,15 +90,10 @@ class LLSSimpleMaskCreate: height=height, radius=radius, ) - combined_mask = combine_masks(created_mask, input_mask, image, combine_mode) + combined_mask = combine_masks(created_mask, input_mask, reference_image, combine_mode) final_mask = apply_invert(combined_mask, bool(invert_mask)) final_mask = apply_softening(final_mask, feather=float(feather), blur=float(blur)) - preview_image = build_preview_image( - image, - final_mask, - overlay_alpha=float(overlay_alpha), - overlay_color=str(overlay_color or "red"), - ) + mask_image = mask_to_image(final_mask) area_info = build_area_info( final_mask, geometry_info, @@ -94,7 +102,7 @@ class LLSSimpleMaskCreate: invert_mask=bool(invert_mask), combine_mode=str(combine_mode or "replace"), ) - return image, final_mask, preview_image, area_info + return final_mask, mask_image, area_info NODE_CLASS_MAPPINGS = {"LLSSimpleMaskCreate": LLSSimpleMaskCreate} diff --git a/mask/mask_preview.py b/mask/mask_preview.py new file mode 100644 index 0000000..a0010f3 --- /dev/null +++ b/mask/mask_preview.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +from .mask_utils import OVERLAY_COLOR_CHOICES, build_preview_image + + +class LLSSimpleMaskPreview: + CATEGORY = "LLS/Mask" + FUNCTION = "preview_mask" + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("preview_image",) + DESCRIPTION = "Overlay a mask on top of an image for preview only." + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "mask": ("MASK",), + "overlay_alpha": ("FLOAT", {"default": 0.4, "min": 0.0, "max": 1.0, "step": 0.01}), + "overlay_color": (OVERLAY_COLOR_CHOICES, {"default": "red"}), + } + } + + def preview_mask(self, image, mask, overlay_alpha, overlay_color): + if image is None: + raise RuntimeError("[LLS] Missing image for LLS Simple Mask Preview.") + if mask is None: + raise RuntimeError("[LLS] Missing mask for LLS Simple Mask Preview.") + preview_image = build_preview_image( + image, + mask, + overlay_alpha=float(overlay_alpha), + overlay_color=str(overlay_color or "red"), + ) + return (preview_image,) + + +NODE_CLASS_MAPPINGS = {"LLSSimpleMaskPreview": LLSSimpleMaskPreview} +NODE_DISPLAY_NAME_MAPPINGS = {"LLSSimpleMaskPreview": "LLS Simple Mask Preview"} diff --git a/mask/mask_utils.py b/mask/mask_utils.py index 72628d6..1f61931 100644 --- a/mask/mask_utils.py +++ b/mask/mask_utils.py @@ -26,6 +26,35 @@ _OVERLAY_COLORS = { } +def create_reference_image( + image_width: int, + image_height: int, + *, + batch: int = 1, + dtype=None, + device=None, +): + _require_torch("[LLS] torch is required for mask generation.") + + width = max(1, int(image_width)) + height = max(1, int(image_height)) + batch = max(1, int(batch)) + kwargs = {"dtype": dtype or torch.float32} + if device is not None: + kwargs["device"] = device + return torch.zeros((batch, height, width, 3), **kwargs) + + +def mask_to_image(mask): + _require_torch("[LLS] torch is required for mask image generation.") + + shape = getattr(mask, "shape", None) + if not isinstance(shape, (list, tuple)) or len(shape) != 3: + raise RuntimeError("[LLS] mask must have shape [batch, height, width].") + normalized = mask.clamp(0.0, 1.0) + return normalized.unsqueeze(-1).repeat(1, 1, 1, 3) + + def get_image_size(image) -> tuple[int, int]: shape = getattr(image, "shape", None) if not isinstance(shape, (list, tuple)) or len(shape) != 4: diff --git a/mask/nodes.py b/mask/nodes.py index caed6c9..141df97 100644 --- a/mask/nodes.py +++ b/mask/nodes.py @@ -3,6 +3,21 @@ LLS / Mask ========== 功能域：遮罩（Mask）操作（对应功能分类总览第 6 节） """ -from .mask_create import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS +from .mask_create import ( + NODE_CLASS_MAPPINGS as MASK_CREATE_NODE_CLASS_MAPPINGS, + NODE_DISPLAY_NAME_MAPPINGS as MASK_CREATE_NODE_DISPLAY_NAME_MAPPINGS, +) +from .mask_preview import ( + NODE_CLASS_MAPPINGS as MASK_PREVIEW_NODE_CLASS_MAPPINGS, + NODE_DISPLAY_NAME_MAPPINGS as MASK_PREVIEW_NODE_DISPLAY_NAME_MAPPINGS, +) + +NODE_CLASS_MAPPINGS = {} +NODE_CLASS_MAPPINGS.update(MASK_CREATE_NODE_CLASS_MAPPINGS) +NODE_CLASS_MAPPINGS.update(MASK_PREVIEW_NODE_CLASS_MAPPINGS) + +NODE_DISPLAY_NAME_MAPPINGS = {} +NODE_DISPLAY_NAME_MAPPINGS.update(MASK_CREATE_NODE_DISPLAY_NAME_MAPPINGS) +NODE_DISPLAY_NAME_MAPPINGS.update(MASK_PREVIEW_NODE_DISPLAY_NAME_MAPPINGS) __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/tests/test_mask_create_docs.py b/tests/test_mask_create_docs.py index 54da737..7ca5f96 100644 --- a/tests/test_mask_create_docs.py +++ b/tests/test_mask_create_docs.py @@ -11,11 +11,15 @@ class TestMaskCreateDocs(unittest.TestCase): readme = (ROOT / "README.md").read_text(encoding="utf-8") self.assertIn("LLS Simple Mask Create", readme) + self.assertIn("LLS Simple Mask Preview", readme) + self.assertIn("image_width", readme) + self.assertIn("image_height", readme) self.assertIn("shape_type", readme) self.assertIn("coordinate_mode", readme) self.assertIn("area_info", readme) - self.assertIn("Load Image -> LLS Simple Mask Create -> LLS Simple Repair Prepare", readme) - self.assertIn("Load Image -> LLS Simple Mask Create -> LLS Simple Mask Draw", readme) + self.assertIn("LLS Simple Mask Create.mask -> LLS Simple Repair Prepare.mask", readme) + self.assertIn("LLS Simple Mask Create.mask_image -> Preview Image", readme) + self.assertIn("Load Image.image + LLS Simple Mask Create.mask -> LLS Simple Mask Preview", readme) def test_readme_lists_supported_shapes_and_combine_modes(self): readme = (ROOT / "README.md").read_text(encoding="utf-8") diff --git a/tests/test_mask_create_node.py b/tests/test_mask_create_node.py index 0ac307d..34a926b 100644 --- a/tests/test_mask_create_node.py +++ b/tests/test_mask_create_node.py @@ -49,7 +49,8 @@ class TestMaskCreateNode(unittest.TestCase): def _create(self, **overrides): params = { - "image": make_image(width=100, height=80, color=0.2), + "image_width": 100, + "image_height": 80, "shape_type": "rectangle", "coordinate_mode": "percent", "center_x": 0.5, @@ -61,19 +62,16 @@ class TestMaskCreateNode(unittest.TestCase): "blur": 0.0, "invert_mask": False, "combine_mode": "replace", - "overlay_alpha": 0.4, - "overlay_color": "red", "input_mask": None, } params.update(overrides) return self.node.create_mask(**params) def test_rectangle_percent_mode_creates_expected_bbox_and_area_info(self): - image_out, mask, preview_image, area_info = self._create() + mask, mask_image, area_info = self._create() self.assertEqual(tuple(mask.shape), (1, 80, 100)) - self.assertEqual(tuple(preview_image.shape), (1, 80, 100, 3)) - self.assertEqual(tuple(image_out.shape), (1, 80, 100, 3)) + self.assertEqual(tuple(mask_image.shape), (1, 80, 100, 3)) self.assertEqual(area_info["image_size"], [100, 80]) self.assertEqual(area_info["shape_type"], "rectangle") self.assertEqual(area_info["coordinate_mode"], "percent") @@ -86,11 +84,13 @@ class TestMaskCreateNode(unittest.TestCase): self.assertAlmostEqual(area_info["area_ratio"], 600.0 / 8000.0, places=4) self.assertEqual(float(mask[0, 40, 50].item()), 1.0) self.assertEqual(float(mask[0, 5, 5].item()), 0.0) + self.assertEqual(float(mask_image[0, 40, 50, 0].item()), 1.0) + self.assertEqual(float(mask_image[0, 5, 5, 0].item()), 0.0) def test_square_pixel_mode_uses_width_as_side(self): - image = make_image(width=64, height=64, color=0.2) - _image_out, mask, _preview_image, area_info = self.node.create_mask( - image=image, + mask, _mask_image, area_info = self.node.create_mask( + image_width=64, + image_height=64, shape_type="square", coordinate_mode="pixel", center_x=20.0, @@ -102,8 +102,6 @@ class TestMaskCreateNode(unittest.TestCase): blur=0.0, invert_mask=False, combine_mode="replace", - overlay_alpha=0.4, - overlay_color="red", input_mask=None, ) @@ -114,9 +112,9 @@ class TestMaskCreateNode(unittest.TestCase): self.assertEqual(float(mask[0, 5, 5].item()), 0.0) def test_circle_generates_centered_mask(self): - image = make_image(width=101, height=101, color=0.2) - _image_out, mask, _preview_image, area_info = self.node.create_mask( - image=image, + mask, _mask_image, area_info = self.node.create_mask( + image_width=101, + image_height=101, shape_type="circle", coordinate_mode="pixel", center_x=50.0, @@ -128,8 +126,6 @@ class TestMaskCreateNode(unittest.TestCase): blur=0.0, invert_mask=False, combine_mode="replace", - overlay_alpha=0.4, - overlay_color="red", input_mask=None, ) @@ -140,9 +136,9 @@ class TestMaskCreateNode(unittest.TestCase): self.assertEqual(float(mask[0, 10, 10].item()), 0.0) def test_ellipse_generates_nonempty_mask(self): - image = make_image(width=120, height=90, color=0.2) - _image_out, mask, _preview_image, area_info = self.node.create_mask( - image=image, + mask, _mask_image, area_info = self.node.create_mask( + image_width=120, + image_height=90, shape_type="ellipse", coordinate_mode="pixel", center_x=60.0, @@ -154,8 +150,6 @@ class TestMaskCreateNode(unittest.TestCase): blur=0.0, invert_mask=False, combine_mode="replace", - overlay_alpha=0.4, - overlay_color="red", input_mask=None, ) @@ -165,36 +159,12 @@ class TestMaskCreateNode(unittest.TestCase): self.assertEqual(float(mask[0, 45, 60].item()), 1.0) self.assertEqual(float(mask[0, 5, 5].item()), 0.0) - def test_overlay_image_uses_requested_color(self): - image = make_image(width=20, height=20, color=0.2) - _image_out, _mask, preview_image, _area_info = self.node.create_mask( - image=image, - shape_type="rectangle", - coordinate_mode="pixel", - center_x=10.0, - center_y=10.0, - width=10.0, - height=10.0, - radius=5.0, - feather=0.0, - blur=0.0, - invert_mask=False, - combine_mode="replace", - overlay_alpha=0.5, - overlay_color="green", - input_mask=None, - ) - - center_pixel = preview_image[0, 10, 10] - self.assertGreater(float(center_pixel[1].item()), float(center_pixel[0].item())) - self.assertGreater(float(center_pixel[1].item()), float(center_pixel[2].item())) - def test_union_subtract_and_intersect_with_input_mask(self): - image = make_image(width=32, height=32, color=0.2) input_mask = make_mask(width=32, height=32, value=1.0) - _image_out, union_mask, _preview, union_info = self.node.create_mask( - image=image, + union_mask, _mask_image, union_info = self.node.create_mask( + image_width=32, + image_height=32, shape_type="rectangle", coordinate_mode="pixel", center_x=16.0, @@ -206,15 +176,14 @@ class TestMaskCreateNode(unittest.TestCase): blur=0.0, invert_mask=False, combine_mode="union", - overlay_alpha=0.4, - overlay_color="red", input_mask=input_mask, ) self.assertEqual(union_info["binary_area_px"], 32 * 32) self.assertTrue(torch.all(union_mask == 1.0)) - _image_out, subtract_mask, _preview, subtract_info = self.node.create_mask( - image=image, + subtract_mask, _mask_image, subtract_info = self.node.create_mask( + image_width=32, + image_height=32, shape_type="rectangle", coordinate_mode="pixel", center_x=16.0, @@ -226,15 +195,14 @@ class TestMaskCreateNode(unittest.TestCase): blur=0.0, invert_mask=False, combine_mode="subtract", - overlay_alpha=0.4, - overlay_color="red", input_mask=input_mask, ) self.assertLess(subtract_info["binary_area_px"], 32 * 32) self.assertEqual(float(subtract_mask[0, 16, 16].item()), 0.0) - _image_out, intersect_mask, _preview, intersect_info = self.node.create_mask( - image=image, + intersect_mask, _mask_image, intersect_info = self.node.create_mask( + image_width=32, + image_height=32, shape_type="rectangle", coordinate_mode="pixel", center_x=16.0, @@ -246,8 +214,6 @@ class TestMaskCreateNode(unittest.TestCase): blur=0.0, invert_mask=False, combine_mode="intersect", - overlay_alpha=0.4, - overlay_color="red", input_mask=input_mask, ) self.assertEqual(intersect_info["binary_area_px"], 100) @@ -255,14 +221,14 @@ class TestMaskCreateNode(unittest.TestCase): self.assertEqual(float(intersect_mask[0, 1, 1].item()), 0.0) def test_invert_mask_flips_final_mask(self): - _image_out, mask, _preview, area_info = self._create(invert_mask=True) + mask, _mask_image, area_info = self._create(invert_mask=True) self.assertEqual(area_info["invert_mask"], True) self.assertEqual(float(mask[0, 40, 50].item()), 0.0) self.assertEqual(float(mask[0, 5, 5].item()), 1.0) def test_feather_and_blur_create_soft_mask(self): - _image_out, mask, _preview, area_info = self._create( + mask, _mask_image, area_info = self._create( shape_type="rectangle", coordinate_mode="pixel", center_x=50.0, @@ -278,7 +244,7 @@ class TestMaskCreateNode(unittest.TestCase): self.assertTrue(bool(((mask > 0.0) & (mask < 1.0)).any().item())) def test_shape_outside_image_returns_black_mask_without_crashing(self): - _image_out, mask, _preview, area_info = self._create( + mask, _mask_image, area_info = self._create( shape_type="circle", coordinate_mode="percent", center_x=2.0, @@ -294,8 +260,9 @@ class TestMaskCreateNode(unittest.TestCase): image = make_image(width=64, height=64, color=0.25) input_mask = make_mask(width=16, height=16, value=1.0) - image_out, mask_out, preview_out, area_info = self.node.create_mask( - image=image, + mask_out, _mask_image, area_info = self.node.create_mask( + image_width=64, + image_height=64, shape_type="rectangle", coordinate_mode="pixel", center_x=32.0, @@ -307,13 +274,11 @@ class TestMaskCreateNode(unittest.TestCase): blur=0.0, invert_mask=False, combine_mode="union", - overlay_alpha=0.4, - overlay_color="red", input_mask=input_mask, ) latent, work_image, work_mask, repair_info, recommended_denoise, positive, negative = self.repair_prepare_node.prepare( - image=image_out, + image=image, mask=mask_out, vae=_FakeVAE(), repair_scope="region", @@ -339,7 +304,6 @@ class TestMaskCreateNode(unittest.TestCase): negative=None, ) - self.assertEqual(tuple(preview_out.shape), (1, 64, 64, 3)) self.assertEqual(tuple(mask_out.shape), (1, 64, 64)) self.assertEqual(tuple(work_image.shape), (1, 64, 64, 3)) self.assertEqual(tuple(work_mask.shape), (1, 64, 64)) diff --git a/tests/test_mask_create_registration.py b/tests/test_mask_create_registration.py index d67e58a..aba3eca 100644 --- a/tests/test_mask_create_registration.py +++ b/tests/test_mask_create_registration.py @@ -23,13 +23,16 @@ class TestMaskCreateRegistration(unittest.TestCase): self.assertEqual(node_cls.CATEGORY, "LLS/Mask") self.assertEqual(node_cls.FUNCTION, "create_mask") - self.assertEqual(node_cls.RETURN_TYPES, ("IMAGE", "MASK", "IMAGE", "LLS_MASK_INFO")) - self.assertEqual(node_cls.RETURN_NAMES, ("image", "mask", "preview_image", "area_info")) + self.assertEqual(node_cls.RETURN_TYPES, ("MASK", "IMAGE", "LLS_MASK_INFO")) + self.assertEqual(node_cls.RETURN_NAMES, ("mask", "mask_image", "area_info")) required = schema["required"] optional = schema["optional"] - self.assertEqual(required["image"], ("IMAGE",)) + self.assertEqual(required["image_width"][0], "INT") + self.assertEqual(required["image_height"][0], "INT") + self.assertEqual(required["image_width"][1]["default"], 1024) + self.assertEqual(required["image_height"][1]["default"], 1024) self.assertEqual(optional["input_mask"], ("MASK",)) self.assertEqual(required["shape_type"][0], ["rectangle", "square", "circle", "ellipse"]) self.assertEqual(required["coordinate_mode"][0], ["pixel", "percent"]) @@ -46,10 +49,35 @@ class TestMaskCreateRegistration(unittest.TestCase): self.assertEqual(required["invert_mask"][0], "BOOLEAN") self.assertEqual(required["combine_mode"][0], ["replace", "union", "subtract", "intersect"]) self.assertEqual(required["combine_mode"][1]["default"], "replace") + self.assertTrue(callable(getattr(node_cls, "create_mask", None))) + + def test_plugin_registers_mask_preview_node(self): + plugin = load_plugin_package() + + self.assertIn("LLSSimpleMaskPreview", plugin.NODE_CLASS_MAPPINGS) + self.assertEqual( + plugin.NODE_DISPLAY_NAME_MAPPINGS["LLSSimpleMaskPreview"], + "LLS Simple Mask Preview", + ) + + def test_mask_preview_schema_matches_contract(self): + plugin = load_plugin_package() + node_cls = plugin.NODE_CLASS_MAPPINGS["LLSSimpleMaskPreview"] + schema = node_cls.INPUT_TYPES() + + self.assertEqual(node_cls.CATEGORY, "LLS/Mask") + self.assertEqual(node_cls.FUNCTION, "preview_mask") + self.assertEqual(node_cls.RETURN_TYPES, ("IMAGE",)) + self.assertEqual(node_cls.RETURN_NAMES, ("preview_image",)) + + required = schema["required"] + + self.assertEqual(required["image"], ("IMAGE",)) + self.assertEqual(required["mask"], ("MASK",)) self.assertEqual(required["overlay_alpha"][0], "FLOAT") self.assertEqual(required["overlay_alpha"][1]["default"], 0.4) self.assertEqual(required["overlay_color"][0], ["red", "green", "blue"]) - self.assertTrue(callable(getattr(node_cls, "create_mask", None))) + self.assertTrue(callable(getattr(node_cls, "preview_mask", None))) if __name__ == "__main__": diff --git a/tests/test_mask_preview_node.py b/tests/test_mask_preview_node.py new file mode 100644 index 0000000..8c4b38b --- /dev/null +++ b/tests/test_mask_preview_node.py @@ -0,0 +1,73 @@ +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_MASK_PREVIEW_RUNTIME_DEPS = torch is not None +MASK_PREVIEW_RUNTIME_DEPS_MESSAGE = "mask preview node tests require torch" + + +def _require_mask_preview_runtime_deps(): + if not HAS_MASK_PREVIEW_RUNTIME_DEPS: + raise unittest.SkipTest(MASK_PREVIEW_RUNTIME_DEPS_MESSAGE) + + +def make_image(width=8, height=8, color=0.25): + _require_mask_preview_runtime_deps() + return torch.full((1, height, width, 3), float(color), dtype=torch.float32) + + +def make_mask(width=8, height=8, value=0.0): + _require_mask_preview_runtime_deps() + return torch.full((1, height, width), float(value), dtype=torch.float32) + + +@unittest.skipUnless(HAS_MASK_PREVIEW_RUNTIME_DEPS, MASK_PREVIEW_RUNTIME_DEPS_MESSAGE) +class TestMaskPreviewNode(unittest.TestCase): + @classmethod + def setUpClass(cls): + plugin = load_plugin_package() + cls.node = plugin.NODE_CLASS_MAPPINGS["LLSSimpleMaskPreview"]() + + def test_preview_overlays_requested_color(self): + image = make_image(width=20, height=20, color=0.2) + mask = make_mask(width=20, height=20, value=1.0) + + (preview_image,) = self.node.preview_mask( + image=image, + mask=mask, + overlay_alpha=0.5, + overlay_color="green", + ) + + center_pixel = preview_image[0, 10, 10] + self.assertEqual(tuple(preview_image.shape), (1, 20, 20, 3)) + self.assertGreater(float(center_pixel[1].item()), float(center_pixel[0].item())) + self.assertGreater(float(center_pixel[1].item()), float(center_pixel[2].item())) + + def test_preview_resizes_mask_to_match_image(self): + image = make_image(width=20, height=20, color=0.2) + mask = make_mask(width=10, height=10, value=1.0) + + (preview_image,) = self.node.preview_mask( + image=image, + mask=mask, + overlay_alpha=0.4, + overlay_color="red", + ) + + self.assertEqual(tuple(preview_image.shape), (1, 20, 20, 3)) + self.assertGreater(float(preview_image[..., 0].mean().item()), float(image[..., 0].mean().item())) + + +if __name__ == "__main__": + unittest.main()