From b2c98aa65c139defd0d9782bf5626fc1ad90bcbc Mon Sep 17 00:00:00 2001 From: Gin3601 Date: Mon, 25 May 2026 13:58:50 +0800 Subject: [PATCH] =?UTF-8?q?=E5=88=B6=E4=BD=9CMSAK=E5=A4=84=E7=90=86=20?= =?UTF-8?q?=E8=8A=82=E7=82=B9?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 2 +- mask/mask_combine.py | 49 +++++ mask/mask_create.py | 6 +- mask/mask_process.py | 99 ++++++++++ mask/nodes.py | 12 ++ mask/utils.py | 253 ++++++++++++++++++++++++ tests/test_mask_create_node.py | 22 +-- tests/test_mask_create_registration.py | 4 +- tests/test_mask_process_node.py | 208 +++++++++++++++++++ tests/test_mask_process_registration.py | 85 ++++++++ 10 files changed, 723 insertions(+), 17 deletions(-) create mode 100644 mask/mask_combine.py create mode 100644 mask/mask_process.py create mode 100644 mask/utils.py create mode 100644 tests/test_mask_process_node.py create mode 100644 tests/test_mask_process_registration.py diff --git a/README.md b/README.md index c764ddb..55276d3 100644 --- a/README.md +++ b/README.md @@ -136,8 +136,8 @@ node/ | 名称 | 类型 | 说明 | |------|------|------| -| `mask` | `MASK` | 最终遮罩，白色 / `1` 表示要重绘，黑色 / `0` 表示不重绘 | | `mask_image` | `IMAGE` | 将 `mask` 直接可视化为黑白图像，方便单独预览 | +| `mask` | `MASK` | 最终遮罩，白色 / `1` 表示要重绘，黑色 / `0` 表示不重绘 | | `area_info` | `LLS_MASK_INFO` | 包含面积、占比、bbox、坐标模式等信息的 dict | `area_info` 主要包含： diff --git a/mask/mask_combine.py b/mask/mask_combine.py new file mode 100644 index 0000000..67ffc35 --- /dev/null +++ b/mask/mask_combine.py @@ -0,0 +1,49 @@ +from __future__ import annotations + +from .utils import align_mask_batch, normalize_mask, resize_mask_to_hw, torch + + +COMBINE_MODE_CHOICES = ["add", "subtract", "intersect", "xor", "max", "min"] + + +class LLSMaskCombine: + CATEGORY = "LLS/Mask" + FUNCTION = "combine" + RETURN_TYPES = ("MASK",) + RETURN_NAMES = ("mask",) + DESCRIPTION = "Combine two masks with logical-style operations." + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mask_a": ("MASK",), + "mask_b": ("MASK",), + "mode": (COMBINE_MODE_CHOICES, {"default": "add"}), + } + } + + def combine(self, mask_a, mask_b, mode): + mask_a = normalize_mask(mask_a) + mask_b = align_mask_batch( + resize_mask_to_hw(normalize_mask(mask_b), int(mask_a.shape[1]), int(mask_a.shape[2])), + int(mask_a.shape[0]), + ) + mode = str(mode or "add") + + if mode in {"add", "max"}: + combined = torch.maximum(mask_a, mask_b) + elif mode == "subtract": + combined = torch.clamp(mask_a - mask_b, 0.0, 1.0) + elif mode in {"intersect", "min"}: + combined = torch.minimum(mask_a, mask_b) + elif mode == "xor": + combined = torch.abs(mask_a - mask_b) + else: + raise RuntimeError(f"[LLS] Unsupported mask combine mode '{mode}'.") + + return (combined.clamp(0.0, 1.0).to(dtype=mask_a.dtype),) + + +NODE_CLASS_MAPPINGS = {"LLSMaskCombine": LLSMaskCombine} +NODE_DISPLAY_NAME_MAPPINGS = {"LLSMaskCombine": "LLS Mask Combine"} diff --git a/mask/mask_create.py b/mask/mask_create.py index a130835..9d63cbb 100644 --- a/mask/mask_create.py +++ b/mask/mask_create.py @@ -27,8 +27,8 @@ def _resolve_batch_for_mask_creation(input_mask) -> int: class LLSSimpleMaskCreate: CATEGORY = "LLS/Mask" FUNCTION = "create_mask" - RETURN_TYPES = ("MASK", "IMAGE", MASK_INFO_TYPE) - RETURN_NAMES = ("mask", "mask_image", "area_info") + RETURN_TYPES = ("IMAGE", "MASK", MASK_INFO_TYPE) + RETURN_NAMES = ("mask_image", "mask", "area_info") DESCRIPTION = "Create a geometric mask image from width, height, and shape parameters." @classmethod @@ -102,7 +102,7 @@ class LLSSimpleMaskCreate: invert_mask=bool(invert_mask), combine_mode=str(combine_mode or "replace"), ) - return final_mask, mask_image, area_info + return mask_image, final_mask, area_info NODE_CLASS_MAPPINGS = {"LLSSimpleMaskCreate": LLSSimpleMaskCreate} diff --git a/mask/mask_process.py b/mask/mask_process.py new file mode 100644 index 0000000..5bddebe --- /dev/null +++ b/mask/mask_process.py @@ -0,0 +1,99 @@ +from __future__ import annotations + +from .utils import ( + apply_blur, + apply_grow, + apply_shrink, + align_mask_batch, + fill_mask_holes, + normalize_mask, + remove_small_mask_regions, + resize_mask_to_hw, + resolve_image_size, +) + + +PROCESS_OPERATION_CHOICES = [ + "passthrough", + "threshold", + "invert", + "grow", + "shrink", + "blur", + "feather", + "fill_holes", + "remove_small_regions", + "smooth", + "clamp", + "resize_to_image", +] + + +class LLSMaskProcess: + CATEGORY = "LLS/Mask" + FUNCTION = "process" + RETURN_TYPES = ("MASK",) + RETURN_NAMES = ("mask",) + DESCRIPTION = "Apply common processing operations to a single mask." + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mask": ("MASK",), + "operation": (PROCESS_OPERATION_CHOICES, {"default": "passthrough"}), + "value_float": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + "value_int": ("INT", {"default": 8, "min": -512, "max": 512, "step": 1}), + }, + "optional": { + "image": ("IMAGE",), + }, + } + + def process(self, mask, operation, value_float, value_int, image=None): + normalized = normalize_mask(mask) + operation = str(operation or "passthrough") + + if operation == "passthrough": + processed = normalized + elif operation == "threshold": + processed = (normalized >= float(value_float)).to(dtype=normalized.dtype) + elif operation == "invert": + processed = 1.0 - normalized + elif operation == "grow": + processed = apply_grow(normalized, value_int) + elif operation == "shrink": + processed = apply_shrink(normalized, value_int) + elif operation == "blur": + processed = apply_blur(normalized, value_int) + elif operation == "feather": + processed = apply_blur(normalized, value_int) + elif operation == "fill_holes": + try: + processed = fill_mask_holes(normalized) + except Exception: + processed = normalized + elif operation == "remove_small_regions": + try: + processed = remove_small_mask_regions(normalized, value_int) + except Exception: + processed = normalized + elif operation == "smooth": + processed = apply_blur(normalized, max(1, int(value_int) if int(value_int) > 0 else 1)) + elif operation == "clamp": + processed = normalized + elif operation == "resize_to_image": + image_size = resolve_image_size(image) + if image_size is None: + processed = normalized + else: + batch, height, width = image_size + processed = align_mask_batch(resize_mask_to_hw(normalized, height, width), batch) + else: + raise RuntimeError(f"[LLS] Unsupported mask operation '{operation}'.") + + return (processed.clamp(0.0, 1.0).to(dtype=normalized.dtype),) + + +NODE_CLASS_MAPPINGS = {"LLSMaskProcess": LLSMaskProcess} +NODE_DISPLAY_NAME_MAPPINGS = {"LLSMaskProcess": "LLS Mask Process"} diff --git a/mask/nodes.py b/mask/nodes.py index 35f56f6..f522bc5 100644 --- a/mask/nodes.py +++ b/mask/nodes.py @@ -7,11 +7,23 @@ 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_process import ( + NODE_CLASS_MAPPINGS as MASK_PROCESS_NODE_CLASS_MAPPINGS, + NODE_DISPLAY_NAME_MAPPINGS as MASK_PROCESS_NODE_DISPLAY_NAME_MAPPINGS, +) +from .mask_combine import ( + NODE_CLASS_MAPPINGS as MASK_COMBINE_NODE_CLASS_MAPPINGS, + NODE_DISPLAY_NAME_MAPPINGS as MASK_COMBINE_NODE_DISPLAY_NAME_MAPPINGS, +) NODE_CLASS_MAPPINGS = {} NODE_CLASS_MAPPINGS.update(MASK_CREATE_NODE_CLASS_MAPPINGS) +NODE_CLASS_MAPPINGS.update(MASK_PROCESS_NODE_CLASS_MAPPINGS) +NODE_CLASS_MAPPINGS.update(MASK_COMBINE_NODE_CLASS_MAPPINGS) NODE_DISPLAY_NAME_MAPPINGS = {} NODE_DISPLAY_NAME_MAPPINGS.update(MASK_CREATE_NODE_DISPLAY_NAME_MAPPINGS) +NODE_DISPLAY_NAME_MAPPINGS.update(MASK_PROCESS_NODE_DISPLAY_NAME_MAPPINGS) +NODE_DISPLAY_NAME_MAPPINGS.update(MASK_COMBINE_NODE_DISPLAY_NAME_MAPPINGS) __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/mask/utils.py b/mask/utils.py new file mode 100644 index 0000000..d15ca2c --- /dev/null +++ b/mask/utils.py @@ -0,0 +1,253 @@ +from __future__ import annotations + +from collections import deque +import math + +try: + import numpy as np +except Exception as exc: # pragma: no cover - optional runtime dependency + np = None + _NUMPY_ERR = exc +else: # pragma: no cover - exercised in runtime/tests + _NUMPY_ERR = None + +try: + import torch + import torch.nn.functional as torch_nn_functional +except Exception as exc: # pragma: no cover - optional runtime dependency + torch = None + torch_nn_functional = None + _TORCH_ERR = exc +else: # pragma: no cover - exercised in runtime/tests + _TORCH_ERR = None + + +def _require_torch(): + if torch is None or torch_nn_functional is None: + raise RuntimeError("[LLS] torch is required for mask processing nodes.") from _TORCH_ERR + + +def _coerce_non_negative_int(value) -> int: + try: + return max(0, int(value)) + except Exception: + return 0 + + +def normalize_mask(mask): + _require_torch() + + if mask is None: + raise RuntimeError("[LLS] mask input is required.") + if not isinstance(mask, torch.Tensor): + mask = torch.as_tensor(mask) + + ndim = int(mask.ndim) + if ndim == 2: + normalized = mask.unsqueeze(0) + elif ndim == 3: + normalized = mask + elif ndim == 4: + if int(mask.shape[-1]) == 1: + normalized = mask.squeeze(-1) + elif int(mask.shape[1]) == 1: + normalized = mask.squeeze(1) + else: + raise RuntimeError( + "[LLS] mask must have shape [H,W], [B,H,W], [B,H,W,1], or [B,1,H,W]." + ) + else: + raise RuntimeError( + "[LLS] mask must have shape [H,W], [B,H,W], [B,H,W,1], or [B,1,H,W]." + ) + + return normalized.to(dtype=torch.float32).clamp(0.0, 1.0) + + +def resize_mask_to_hw(mask, height, width): + normalized = normalize_mask(mask) + target_height = max(1, int(height)) + target_width = max(1, int(width)) + if tuple(normalized.shape[1:]) == (target_height, target_width): + return normalized + + padded = normalized.unsqueeze(1) + resized = torch_nn_functional.interpolate( + padded, + size=(target_height, target_width), + mode="bilinear", + align_corners=False, + ) + return resized.squeeze(1).clamp(0.0, 1.0).to(dtype=torch.float32) + + +def align_mask_batch(mask, target_batch): + normalized = normalize_mask(mask) + batch = int(normalized.shape[0]) + target = max(1, int(target_batch)) + if batch == target: + return normalized + if batch == 1: + return normalized.repeat(target, 1, 1) + if batch < target: + repeats = int(math.ceil(target / batch)) + return normalized.repeat(repeats, 1, 1)[:target] + return normalized[:target] + + +def apply_grow(mask, radius): + normalized = normalize_mask(mask) + grow_radius = _coerce_non_negative_int(radius) + if grow_radius == 0: + return normalized + + kernel_size = (grow_radius * 2) + 1 + grown = torch_nn_functional.max_pool2d( + normalized.unsqueeze(1), + kernel_size=kernel_size, + stride=1, + padding=grow_radius, + ) + return grown.squeeze(1).clamp(0.0, 1.0).to(dtype=torch.float32) + + +def apply_shrink(mask, radius): + normalized = normalize_mask(mask) + shrink_radius = _coerce_non_negative_int(radius) + if shrink_radius == 0: + return normalized + + inverted = 1.0 - normalized + kernel_size = (shrink_radius * 2) + 1 + pooled = torch_nn_functional.max_pool2d( + inverted.unsqueeze(1), + kernel_size=kernel_size, + stride=1, + padding=shrink_radius, + ) + return (1.0 - pooled.squeeze(1)).clamp(0.0, 1.0).to(dtype=torch.float32) + + +def apply_blur(mask, radius): + normalized = normalize_mask(mask) + blur_radius = _coerce_non_negative_int(radius) + if blur_radius == 0: + return normalized + + kernel_size = (blur_radius * 2) + 1 + padded = torch_nn_functional.pad( + normalized.unsqueeze(1), + (blur_radius, blur_radius, blur_radius, blur_radius), + mode="replicate", + ) + blurred = torch_nn_functional.avg_pool2d( + padded, + kernel_size=kernel_size, + stride=1, + ) + return blurred.squeeze(1).clamp(0.0, 1.0).to(dtype=torch.float32) + + +def fill_mask_holes(mask): + normalized = normalize_mask(mask) + if np is None: + return normalized + + binary = (normalized > 0.5).detach().cpu().numpy().astype(np.uint8) + filled = np.zeros_like(binary, dtype=np.float32) + for index, sample in enumerate(binary): + filled[index] = _fill_holes_binary(sample) + + return torch.from_numpy(filled).to(device=normalized.device, dtype=torch.float32).clamp(0.0, 1.0) + + +def remove_small_mask_regions(mask, min_area): + normalized = normalize_mask(mask) + area_threshold = _coerce_non_negative_int(min_area) + if area_threshold <= 1 or np is None: + return normalized + + binary = (normalized > 0.5).detach().cpu().numpy().astype(np.uint8) + cleaned = np.zeros_like(binary, dtype=np.float32) + for index, sample in enumerate(binary): + cleaned[index] = _remove_small_regions_binary(sample, area_threshold) + + return torch.from_numpy(cleaned).to(device=normalized.device, dtype=torch.float32).clamp(0.0, 1.0) + + +def resolve_image_size(image): + shape = tuple(getattr(image, "shape", ())) + if len(shape) != 4: + return None + batch = max(1, int(shape[0])) + height = max(1, int(shape[1])) + width = max(1, int(shape[2])) + return batch, height, width + + +def _fill_holes_binary(sample): + height, width = sample.shape + visited = np.zeros((height, width), dtype=bool) + queue = deque() + + def enqueue(y, x): + if 0 <= y < height and 0 <= x < width and sample[y, x] == 0 and not visited[y, x]: + visited[y, x] = True + queue.append((y, x)) + + for x in range(width): + enqueue(0, x) + enqueue(height - 1, x) + for y in range(height): + enqueue(y, 0) + enqueue(y, width - 1) + + while queue: + y, x = queue.popleft() + enqueue(y - 1, x) + enqueue(y + 1, x) + enqueue(y, x - 1) + enqueue(y, x + 1) + + result = sample.copy() + result[(sample == 0) & (~visited)] = 1 + return result.astype(np.float32) + + +def _remove_small_regions_binary(sample, min_area): + height, width = sample.shape + visited = np.zeros((height, width), dtype=bool) + result = np.zeros((height, width), dtype=np.float32) + + for start_y in range(height): + for start_x in range(width): + if sample[start_y, start_x] == 0 or visited[start_y, start_x]: + continue + + queue = deque([(start_y, start_x)]) + visited[start_y, start_x] = True + region = [] + + while queue: + y, x = queue.popleft() + region.append((y, x)) + for next_y, next_x in ( + (y - 1, x), + (y + 1, x), + (y, x - 1), + (y, x + 1), + ): + if ( + 0 <= next_y < height + and 0 <= next_x < width + and sample[next_y, next_x] == 1 + and not visited[next_y, next_x] + ): + visited[next_y, next_x] = True + queue.append((next_y, next_x)) + + if len(region) >= min_area: + for y, x in region: + result[y, x] = 1.0 + + return result diff --git a/tests/test_mask_create_node.py b/tests/test_mask_create_node.py index 34a926b..b71ed4b 100644 --- a/tests/test_mask_create_node.py +++ b/tests/test_mask_create_node.py @@ -68,7 +68,7 @@ class TestMaskCreateNode(unittest.TestCase): return self.node.create_mask(**params) def test_rectangle_percent_mode_creates_expected_bbox_and_area_info(self): - mask, mask_image, area_info = self._create() + mask_image, mask, area_info = self._create() self.assertEqual(tuple(mask.shape), (1, 80, 100)) self.assertEqual(tuple(mask_image.shape), (1, 80, 100, 3)) @@ -88,7 +88,7 @@ class TestMaskCreateNode(unittest.TestCase): self.assertEqual(float(mask_image[0, 5, 5, 0].item()), 0.0) def test_square_pixel_mode_uses_width_as_side(self): - mask, _mask_image, area_info = self.node.create_mask( + _mask_image, mask, area_info = self.node.create_mask( image_width=64, image_height=64, shape_type="square", @@ -112,7 +112,7 @@ class TestMaskCreateNode(unittest.TestCase): self.assertEqual(float(mask[0, 5, 5].item()), 0.0) def test_circle_generates_centered_mask(self): - mask, _mask_image, area_info = self.node.create_mask( + _mask_image, mask, area_info = self.node.create_mask( image_width=101, image_height=101, shape_type="circle", @@ -136,7 +136,7 @@ class TestMaskCreateNode(unittest.TestCase): self.assertEqual(float(mask[0, 10, 10].item()), 0.0) def test_ellipse_generates_nonempty_mask(self): - mask, _mask_image, area_info = self.node.create_mask( + _mask_image, mask, area_info = self.node.create_mask( image_width=120, image_height=90, shape_type="ellipse", @@ -162,7 +162,7 @@ class TestMaskCreateNode(unittest.TestCase): def test_union_subtract_and_intersect_with_input_mask(self): input_mask = make_mask(width=32, height=32, value=1.0) - union_mask, _mask_image, union_info = self.node.create_mask( + _mask_image, union_mask, union_info = self.node.create_mask( image_width=32, image_height=32, shape_type="rectangle", @@ -181,7 +181,7 @@ class TestMaskCreateNode(unittest.TestCase): self.assertEqual(union_info["binary_area_px"], 32 * 32) self.assertTrue(torch.all(union_mask == 1.0)) - subtract_mask, _mask_image, subtract_info = self.node.create_mask( + _mask_image, subtract_mask, subtract_info = self.node.create_mask( image_width=32, image_height=32, shape_type="rectangle", @@ -200,7 +200,7 @@ class TestMaskCreateNode(unittest.TestCase): self.assertLess(subtract_info["binary_area_px"], 32 * 32) self.assertEqual(float(subtract_mask[0, 16, 16].item()), 0.0) - intersect_mask, _mask_image, intersect_info = self.node.create_mask( + _mask_image, intersect_mask, intersect_info = self.node.create_mask( image_width=32, image_height=32, shape_type="rectangle", @@ -221,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): - mask, _mask_image, area_info = self._create(invert_mask=True) + _mask_image, mask, 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): - mask, _mask_image, area_info = self._create( + _mask_image, mask, area_info = self._create( shape_type="rectangle", coordinate_mode="pixel", center_x=50.0, @@ -244,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): - mask, _mask_image, area_info = self._create( + _mask_image, mask, area_info = self._create( shape_type="circle", coordinate_mode="percent", center_x=2.0, @@ -260,7 +260,7 @@ 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) - mask_out, _mask_image, area_info = self.node.create_mask( + _mask_image, mask_out, area_info = self.node.create_mask( image_width=64, image_height=64, shape_type="rectangle", diff --git a/tests/test_mask_create_registration.py b/tests/test_mask_create_registration.py index 1169333..76df192 100644 --- a/tests/test_mask_create_registration.py +++ b/tests/test_mask_create_registration.py @@ -23,8 +23,8 @@ class TestMaskCreateRegistration(unittest.TestCase): self.assertEqual(node_cls.CATEGORY, "LLS/Mask") self.assertEqual(node_cls.FUNCTION, "create_mask") - self.assertEqual(node_cls.RETURN_TYPES, ("MASK", "IMAGE", "LLS_MASK_INFO")) - self.assertEqual(node_cls.RETURN_NAMES, ("mask", "mask_image", "area_info")) + self.assertEqual(node_cls.RETURN_TYPES, ("IMAGE", "MASK", "LLS_MASK_INFO")) + self.assertEqual(node_cls.RETURN_NAMES, ("mask_image", "mask", "area_info")) required = schema["required"] optional = schema["optional"] diff --git a/tests/test_mask_process_node.py b/tests/test_mask_process_node.py new file mode 100644 index 0000000..0cd71e6 --- /dev/null +++ b/tests/test_mask_process_node.py @@ -0,0 +1,208 @@ +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_PROCESS_RUNTIME_DEPS = torch is not None +MASK_PROCESS_RUNTIME_DEPS_MESSAGE = "mask processing node tests require torch" + + +def _require_mask_process_runtime_deps(): + if not HAS_MASK_PROCESS_RUNTIME_DEPS: + raise unittest.SkipTest(MASK_PROCESS_RUNTIME_DEPS_MESSAGE) + + +@unittest.skipUnless(HAS_MASK_PROCESS_RUNTIME_DEPS, MASK_PROCESS_RUNTIME_DEPS_MESSAGE) +class TestMaskProcessNode(unittest.TestCase): + @classmethod + def setUpClass(cls): + plugin = load_plugin_package() + cls.process_node = plugin.NODE_CLASS_MAPPINGS["LLSMaskProcess"]() + cls.combine_node = plugin.NODE_CLASS_MAPPINGS["LLSMaskCombine"]() + + def assert_is_normalized_mask(self, mask, shape): + self.assertEqual(tuple(mask.shape), shape) + self.assertEqual(mask.dtype, torch.float32) + self.assertGreaterEqual(float(mask.min().item()), 0.0) + self.assertLessEqual(float(mask.max().item()), 1.0) + + def test_passthrough_normalizes_supported_mask_shapes(self): + cases = { + "hw": torch.tensor([[-1.0, 0.4], [1.2, 0.7]], dtype=torch.float64), + "bhw": torch.tensor([[[0.0, 1.0], [1.0, 0.0]]], dtype=torch.float32), + "bhw1": torch.tensor([[[[0.0], [1.0]], [[1.0], [0.0]]]], dtype=torch.float32), + "b1hw": torch.tensor([[[[0.0, 1.0], [1.0, 0.0]]]], dtype=torch.float32), + } + + for name, source in cases.items(): + with self.subTest(name=name): + (mask,) = self.process_node.process( + mask=source, + operation="passthrough", + value_float=0.5, + value_int=8, + ) + self.assert_is_normalized_mask(mask, (1, 2, 2)) + + def test_threshold_and_invert_produce_expected_values(self): + source = torch.tensor([[[0.2, 0.5, 0.7]]], dtype=torch.float32) + + (thresholded,) = self.process_node.process( + mask=source, + operation="threshold", + value_float=0.5, + value_int=8, + ) + (inverted,) = self.process_node.process( + mask=source, + operation="invert", + value_float=0.5, + value_int=8, + ) + + self.assertTrue( + torch.equal( + thresholded, + torch.tensor([[[0.0, 1.0, 1.0]]], dtype=torch.float32), + ) + ) + self.assertTrue( + torch.allclose( + inverted, + torch.tensor([[[0.8, 0.5, 0.3]]], dtype=torch.float32), + ) + ) + + def test_grow_and_shrink_do_not_error_and_keep_normalized_output(self): + source = torch.zeros((1, 5, 5), dtype=torch.float32) + source[0, 2, 2] = 1.0 + + (grown,) = self.process_node.process( + mask=source, + operation="grow", + value_float=0.5, + value_int=1, + ) + (shrunk,) = self.process_node.process( + mask=grown, + operation="shrink", + value_float=0.5, + value_int=1, + ) + + self.assert_is_normalized_mask(grown, (1, 5, 5)) + self.assert_is_normalized_mask(shrunk, (1, 5, 5)) + self.assertGreater(float(grown.sum().item()), float(source.sum().item())) + self.assertLessEqual(float(shrunk.sum().item()), float(grown.sum().item())) + + def test_blur_and_feather_do_not_error_and_create_soft_edges(self): + source = torch.zeros((1, 7, 7), dtype=torch.float32) + source[:, 2:5, 2:5] = 1.0 + + (blurred,) = self.process_node.process( + mask=source, + operation="blur", + value_float=0.5, + value_int=1, + ) + (feathered,) = self.process_node.process( + mask=source, + operation="feather", + value_float=0.5, + value_int=2, + ) + + self.assert_is_normalized_mask(blurred, (1, 7, 7)) + self.assert_is_normalized_mask(feathered, (1, 7, 7)) + self.assertTrue(bool(((blurred > 0.0) & (blurred < 1.0)).any().item())) + self.assertTrue(bool(((feathered > 0.0) & (feathered < 1.0)).any().item())) + + def test_fill_holes_remove_small_regions_smooth_and_clamp_are_stable(self): + with_hole = torch.ones((1, 5, 5), dtype=torch.float32) + with_hole[0, 2, 2] = 0.0 + + (filled,) = self.process_node.process( + mask=with_hole, + operation="fill_holes", + value_float=0.5, + value_int=8, + ) + self.assertEqual(float(filled[0, 2, 2].item()), 1.0) + + with_small_region = torch.zeros((1, 6, 6), dtype=torch.float32) + with_small_region[0, 0, 0] = 1.0 + with_small_region[0, 3:5, 3:5] = 1.0 + (removed,) = self.process_node.process( + mask=with_small_region, + operation="remove_small_regions", + value_float=0.5, + value_int=2, + ) + self.assertEqual(float(removed[0, 0, 0].item()), 0.0) + self.assertEqual(float(removed[0, 3, 3].item()), 1.0) + + jagged = torch.tensor([[[0.0, 1.0, 0.0], [1.0, 1.0, 1.0], [0.0, 1.0, 0.0]]], dtype=torch.float32) + (smoothed,) = self.process_node.process( + mask=jagged, + operation="smooth", + value_float=0.5, + value_int=1, + ) + (clamped,) = self.process_node.process( + mask=torch.tensor([[[-2.0, 2.0]]], dtype=torch.float32), + operation="clamp", + value_float=0.5, + value_int=8, + ) + + self.assert_is_normalized_mask(smoothed, (1, 3, 3)) + self.assert_is_normalized_mask(clamped, (1, 1, 2)) + + def test_resize_to_image_uses_optional_image_dimensions_and_batch(self): + source = torch.ones((1, 2, 3), dtype=torch.float32) + image = torch.zeros((2, 5, 7, 3), dtype=torch.float32) + + (resized,) = self.process_node.process( + mask=source, + operation="resize_to_image", + value_float=0.5, + value_int=8, + image=image, + ) + + self.assert_is_normalized_mask(resized, (2, 5, 7)) + self.assertTrue(torch.allclose(resized, torch.ones((2, 5, 7), dtype=torch.float32))) + + def test_mask_combine_modes_produce_expected_values(self): + mask_a = torch.tensor([[[1.0, 0.5], [0.0, 1.0]]], dtype=torch.float32) + mask_b = torch.tensor([[[0.5, 1.0], [0.0, 0.25]]], dtype=torch.float32) + expectations = { + "add": torch.tensor([[[1.0, 1.0], [0.0, 1.0]]], dtype=torch.float32), + "subtract": torch.tensor([[[0.5, 0.0], [0.0, 0.75]]], dtype=torch.float32), + "intersect": torch.tensor([[[0.5, 0.5], [0.0, 0.25]]], dtype=torch.float32), + "xor": torch.tensor([[[0.5, 0.5], [0.0, 0.75]]], dtype=torch.float32), + "max": torch.tensor([[[1.0, 1.0], [0.0, 1.0]]], dtype=torch.float32), + "min": torch.tensor([[[0.5, 0.5], [0.0, 0.25]]], dtype=torch.float32), + } + + for mode, expected in expectations.items(): + with self.subTest(mode=mode): + (combined,) = self.combine_node.combine(mask_a=mask_a, mask_b=mask_b, mode=mode) + self.assert_is_normalized_mask(combined, (1, 2, 2)) + self.assertTrue(torch.allclose(combined, expected)) + + def test_mask_combine_resizes_and_aligns_batch_to_mask_a(self): + mask_a = torch.full((2, 4, 6), 0.25, dtype=torch.float32) + mask_b = torch.ones((1, 2, 3), dtype=torch.float32) + + (combined,) = self.combine_node.combine(mask_a=mask_a, mask_b=mask_b, mode="add") + + self.assert_is_normalized_mask(combined, (2, 4, 6)) + self.assertTrue(torch.allclose(combined, torch.ones((2, 4, 6), dtype=torch.float32))) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_mask_process_registration.py b/tests/test_mask_process_registration.py new file mode 100644 index 0000000..6f0e53f --- /dev/null +++ b/tests/test_mask_process_registration.py @@ -0,0 +1,85 @@ +import unittest + +try: + from .test_mask_draw_helpers import load_plugin_package +except ImportError: + from test_mask_draw_helpers import load_plugin_package + + +class TestMaskProcessRegistration(unittest.TestCase): + def test_plugin_registers_mask_processing_nodes(self): + plugin = load_plugin_package() + + self.assertIn("LLSMaskProcess", plugin.NODE_CLASS_MAPPINGS) + self.assertIn("LLSMaskCombine", plugin.NODE_CLASS_MAPPINGS) + self.assertEqual( + plugin.NODE_DISPLAY_NAME_MAPPINGS["LLSMaskProcess"], + "LLS Mask Process", + ) + self.assertEqual( + plugin.NODE_DISPLAY_NAME_MAPPINGS["LLSMaskCombine"], + "LLS Mask Combine", + ) + + def test_mask_process_schema_matches_contract(self): + plugin = load_plugin_package() + node_cls = plugin.NODE_CLASS_MAPPINGS["LLSMaskProcess"] + schema = node_cls.INPUT_TYPES() + required = schema["required"] + optional = schema["optional"] + + self.assertEqual(node_cls.CATEGORY, "LLS/Mask") + self.assertEqual(node_cls.FUNCTION, "process") + self.assertEqual(node_cls.RETURN_TYPES, ("MASK",)) + self.assertEqual(node_cls.RETURN_NAMES, ("mask",)) + self.assertEqual(required["mask"], ("MASK",)) + self.assertEqual( + required["operation"][0], + [ + "passthrough", + "threshold", + "invert", + "grow", + "shrink", + "blur", + "feather", + "fill_holes", + "remove_small_regions", + "smooth", + "clamp", + "resize_to_image", + ], + ) + self.assertEqual(required["value_float"][0], "FLOAT") + self.assertEqual(required["value_float"][1]["default"], 0.5) + self.assertEqual(required["value_float"][1]["min"], 0.0) + self.assertEqual(required["value_float"][1]["max"], 1.0) + self.assertEqual(required["value_int"][0], "INT") + self.assertEqual(required["value_int"][1]["default"], 8) + self.assertEqual(required["value_int"][1]["min"], -512) + self.assertEqual(required["value_int"][1]["max"], 512) + self.assertEqual(optional["image"], ("IMAGE",)) + + def test_mask_combine_schema_matches_contract(self): + plugin = load_plugin_package() + node_cls = plugin.NODE_CLASS_MAPPINGS["LLSMaskCombine"] + schema = node_cls.INPUT_TYPES() + required = schema["required"] + + self.assertEqual(node_cls.CATEGORY, "LLS/Mask") + self.assertEqual(node_cls.FUNCTION, "combine") + self.assertEqual(node_cls.RETURN_TYPES, ("MASK",)) + self.assertEqual(node_cls.RETURN_NAMES, ("mask",)) + self.assertEqual(required["mask_a"], ("MASK",)) + self.assertEqual(required["mask_b"], ("MASK",)) + self.assertEqual(required["mode"][0], ["add", "subtract", "intersect", "xor", "max", "min"]) + + def test_plugin_does_not_register_mask_preview_node(self): + plugin = load_plugin_package() + + self.assertNotIn("LLSMaskPreview", plugin.NODE_CLASS_MAPPINGS) + self.assertNotIn("LLSMaskPreview", plugin.NODE_DISPLAY_NAME_MAPPINGS) + + +if __name__ == "__main__": + unittest.main()