66 lines
1.7 KiB
Python
66 lines
1.7 KiB
Python
from __future__ import annotations
|
|
import numpy as np
|
|
from backend.node_registry import register_node
|
|
|
|
|
|
_MASK_BOOLEAN_OPERATIONS = {
|
|
"and": lambda a, b: a & b,
|
|
"or": lambda a, b: a | b,
|
|
"xor": lambda a, b: a ^ b,
|
|
"xnor": lambda a, b: ~(a ^ b),
|
|
"nand": lambda a, b: ~(a & b),
|
|
"nor": lambda a, b: ~(a | b),
|
|
"a_minus_b": lambda a, b: a & ~b,
|
|
"b_minus_a": lambda a, b: b & ~a,
|
|
"a": lambda a, b: a,
|
|
"b": lambda a, b: b,
|
|
"not_a": lambda a, b: ~a,
|
|
"not_b": lambda a, b: ~b,
|
|
"a_implies_b": lambda a, b: ~a | b,
|
|
"b_implies_a": lambda a, b: ~b | a,
|
|
"false": lambda a, b: np.zeros_like(a, dtype=bool),
|
|
"true": lambda a, b: np.ones_like(a, dtype=bool),
|
|
}
|
|
|
|
|
|
@register_node(display_name="Mask Operations")
|
|
class MaskOperations:
|
|
_CUSTOM_PREVIEW = True
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"mask_a": ("IMAGE",),
|
|
"mask_b": ("IMAGE",),
|
|
"operation": (list(_MASK_BOOLEAN_OPERATIONS.keys()),),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("mask",)
|
|
FUNCTION = "process"
|
|
|
|
DESCRIPTION = (
|
|
"Apply boolean logic to two binary masks. Includes AND, OR, XOR, NAND, NOR, "
|
|
"XNOR, directional subtraction, implication, pass-through, and constant true/false outputs."
|
|
)
|
|
|
|
def process(
|
|
self,
|
|
mask_a: np.ndarray,
|
|
mask_b: np.ndarray,
|
|
operation: str,
|
|
) -> tuple:
|
|
a = mask_a > 127
|
|
b = mask_b > 127
|
|
|
|
op = _MASK_BOOLEAN_OPERATIONS.get(operation)
|
|
if op is None:
|
|
raise ValueError(f"Unknown mask operation: {operation}")
|
|
result = op(a, b)
|
|
|
|
out = result.astype(np.uint8) * 255
|
|
|
|
return (out,)
|