72 lines
2.4 KiB
Python
72 lines
2.4 KiB
Python
from __future__ import annotations
|
|
import numpy as np
|
|
from backend.node_registry import register_node
|
|
from backend.data_types import DataField, encode_preview
|
|
from backend.nodes.helpers import _mask_overlay, _mask_structure
|
|
|
|
|
|
@register_node(display_name="Mask Morphology")
|
|
class MaskMorphology:
|
|
"""Morphological operations on binary masks.
|
|
|
|
Equivalent to Gwyddion's mask_morph.c (erode, dilate, open, close).
|
|
"""
|
|
_CUSTOM_PREVIEW = True
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"mask": ("IMAGE",),
|
|
"operation": (["dilate", "erode", "open", "close"],),
|
|
"radius": ("INT", {"default": 1, "min": 1, "max": 50, "step": 1}),
|
|
"shape": (["disk", "square"],),
|
|
},
|
|
"optional": {
|
|
"field": ("DATA_FIELD",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("mask",)
|
|
FUNCTION = "process"
|
|
|
|
DESCRIPTION = (
|
|
"Apply morphological operations to a binary mask. "
|
|
"Dilate expands regions, erode shrinks them, "
|
|
"open (erode then dilate) removes small spots, "
|
|
"close (dilate then erode) fills small holes. "
|
|
"Equivalent to Gwyddion mask_morph."
|
|
)
|
|
|
|
_broadcast_fn = None
|
|
_current_node_id: str = ""
|
|
|
|
def process(self, mask: np.ndarray, operation: str, radius: int, shape: str,
|
|
field: DataField | None = None) -> tuple:
|
|
from scipy.ndimage import binary_closing, binary_dilation, binary_erosion, binary_opening
|
|
|
|
binary = mask > 127
|
|
struct = _mask_structure(radius, shape)
|
|
|
|
if operation == "dilate":
|
|
result = binary_dilation(binary, structure=struct)
|
|
elif operation == "erode":
|
|
result = binary_erosion(binary, structure=struct)
|
|
elif operation == "open":
|
|
result = binary_opening(binary, structure=struct)
|
|
elif operation == "close":
|
|
result = binary_closing(binary, structure=struct)
|
|
else:
|
|
raise ValueError(f"Unknown morphological operation: {operation}")
|
|
|
|
out = result.astype(np.uint8) * 255
|
|
|
|
if field is not None and MaskMorphology._broadcast_fn is not None:
|
|
overlay = _mask_overlay(field, out)
|
|
MaskMorphology._broadcast_fn(
|
|
MaskMorphology._current_node_id, encode_preview(overlay),
|
|
)
|
|
|
|
return (out,)
|