87 lines
3.0 KiB
Python
87 lines
3.0 KiB
Python
"""Mask noisify -- add random perturbation to mask boundaries."""
|
|
|
|
from __future__ import annotations
|
|
import numpy as np
|
|
from backend.node_registry import register_node
|
|
from backend.data_types import DataField
|
|
from backend.nodes.helpers import mask_to_bool, bool_to_mask, emit_mask_preview
|
|
|
|
|
|
@register_node(display_name="Mask Noisify")
|
|
class MaskNoisify:
|
|
"""
|
|
Add random perturbation to mask boundaries.
|
|
"""
|
|
_CUSTOM_PREVIEW = True
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"mask": ("IMAGE",),
|
|
"density": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"direction": (["both", "add", "remove"],),
|
|
"boundaries_only": ("BOOLEAN", {"default": True}),
|
|
"seed": ("INT", {"default": 42, "min": 0, "max": 999999}),
|
|
},
|
|
"optional": {
|
|
"field": ("DATA_FIELD",),
|
|
}
|
|
}
|
|
|
|
OUTPUTS = (
|
|
('IMAGE', 'mask'),
|
|
)
|
|
FUNCTION = "process"
|
|
|
|
DESCRIPTION = (
|
|
"Add random noise to a binary mask by flipping pixels near boundaries. "
|
|
"Control the fraction of affected pixels with density, restrict changes "
|
|
"to boundary pixels, and choose whether to add, remove, or both. "
|
|
"Use a fixed seed for reproducible results."
|
|
)
|
|
|
|
KEYWORDS = ("random", "perturb", "boundary", "roughen", "jitter")
|
|
|
|
def process(self, mask: np.ndarray, density: float, direction: str,
|
|
boundaries_only: bool, seed: int,
|
|
field: DataField | None = None) -> tuple:
|
|
binary = mask_to_bool(mask)
|
|
|
|
# Identify boundary pixels: pixels that differ from at least one neighbour
|
|
if boundaries_only:
|
|
boundary = np.zeros_like(binary)
|
|
for shift_axis, shift_dir in [(0, 1), (0, -1), (1, 1), (1, -1)]:
|
|
shifted = np.roll(binary, shift_dir, axis=shift_axis)
|
|
boundary |= (binary != shifted)
|
|
else:
|
|
boundary = np.ones_like(binary)
|
|
|
|
# Select candidate pixels based on direction
|
|
if direction == "add":
|
|
candidates = boundary & ~binary
|
|
elif direction == "remove":
|
|
candidates = boundary & binary
|
|
else: # "both"
|
|
candidates = boundary
|
|
|
|
# Randomly flip density fraction of candidates
|
|
candidate_indices = np.argwhere(candidates)
|
|
n_candidates = len(candidate_indices)
|
|
|
|
if n_candidates > 0 and density > 0:
|
|
rng = np.random.default_rng(seed)
|
|
n_flip = int(round(density * n_candidates))
|
|
n_flip = max(0, min(n_flip, n_candidates))
|
|
if n_flip > 0:
|
|
chosen = rng.choice(n_candidates, size=n_flip, replace=False)
|
|
for idx in chosen:
|
|
r, c = candidate_indices[idx]
|
|
binary[r, c] = ~binary[r, c]
|
|
|
|
out = bool_to_mask(binary)
|
|
|
|
emit_mask_preview(field, out)
|
|
|
|
return (out,)
|