62 lines
1.7 KiB
Python
62 lines
1.7 KiB
Python
from __future__ import annotations
|
|
import numpy as np
|
|
from backend.node_registry import register_node
|
|
from backend.execution_context import emit_preview
|
|
from backend.data_types import DataField, encode_preview
|
|
from backend.nodes.helpers import _mask_overlay
|
|
|
|
|
|
@register_node(display_name="Mask Combine")
|
|
class MaskCombine:
|
|
_CUSTOM_PREVIEW = True
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"mask_a": ("IMAGE",),
|
|
"mask_b": ("IMAGE",),
|
|
"operation": (["and", "or", "xor", "subtract"],),
|
|
},
|
|
"optional": {
|
|
"field": ("DATA_FIELD",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("mask",)
|
|
FUNCTION = "process"
|
|
|
|
DESCRIPTION = (
|
|
"Combine two binary masks with a boolean operation. "
|
|
"AND keeps overlap, OR merges, XOR keeps non-overlapping regions, "
|
|
"subtract removes mask_b from mask_a."
|
|
)
|
|
|
|
_broadcast_fn = None
|
|
_current_node_id: str = ""
|
|
|
|
def process(self, mask_a: np.ndarray, mask_b: np.ndarray, operation: str,
|
|
field: DataField | None = None) -> tuple:
|
|
a = mask_a > 127
|
|
b = mask_b > 127
|
|
|
|
if operation == "and":
|
|
result = a & b
|
|
elif operation == "or":
|
|
result = a | b
|
|
elif operation == "xor":
|
|
result = a ^ b
|
|
elif operation == "subtract":
|
|
result = a & ~b
|
|
else:
|
|
raise ValueError(f"Unknown mask operation: {operation}")
|
|
|
|
out = result.astype(np.uint8) * 255
|
|
|
|
if field is not None:
|
|
overlay = _mask_overlay(field, out)
|
|
emit_preview(encode_preview(overlay))
|
|
|
|
return (out,)
|