114 lines
3.8 KiB
Python
114 lines
3.8 KiB
Python
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
|
|
from backend.data_types import DataField, datafield_to_uint8, encode_preview
|
|
from backend.execution_context import emit_overlay
|
|
from backend.node_registry import register_node
|
|
from backend.nodes.helpers import bool_to_mask, coerce_physical_square
|
|
|
|
|
|
@register_node(display_name="Rectangular Mask")
|
|
class RectangularMask:
|
|
_CUSTOM_PREVIEW = True
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"field": ("DATA_FIELD",),
|
|
"x1": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 1.0, "step": 0.01, "hidden": True}),
|
|
"y1": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 1.0, "step": 0.01, "hidden": True}),
|
|
"x2": ("FLOAT", {"default": 0.8, "min": 0.0, "max": 1.0, "step": 0.01, "hidden": True}),
|
|
"y2": ("FLOAT", {"default": 0.8, "min": 0.0, "max": 1.0, "step": 0.01, "hidden": True}),
|
|
"square": ("BOOLEAN", {"default": False}),
|
|
"invert": ("BOOLEAN", {"default": False}),
|
|
},
|
|
"optional": {
|
|
"corner_a": ("COORD",),
|
|
"corner_b": ("COORD",),
|
|
},
|
|
}
|
|
|
|
OUTPUTS = (
|
|
('IMAGE', 'mask'),
|
|
)
|
|
FUNCTION = "process"
|
|
|
|
DESCRIPTION = (
|
|
"Create a binary mask covering a rectangular region of a DATA_FIELD, "
|
|
"defined by two draggable corners on the preview. Useful for selecting "
|
|
"a region of interest without cropping the image. When 'square' is on, "
|
|
"the mask is coerced to a physical square (the longer side shrinks to "
|
|
"match the shorter, anchored at the top-left corner). When 'invert' is "
|
|
"on, the mask covers everything outside the rectangle instead. "
|
|
"Incoming COORD inputs can lock either corner."
|
|
)
|
|
|
|
KEYWORDS = ("roi", "region", "rectangle", "square", "box", "selection", "crop mask")
|
|
|
|
def process(
|
|
self,
|
|
field: DataField,
|
|
x1: float,
|
|
y1: float,
|
|
x2: float,
|
|
y2: float,
|
|
square: bool,
|
|
invert: bool,
|
|
corner_a=None,
|
|
corner_b=None,
|
|
) -> tuple:
|
|
if corner_a is not None:
|
|
x1, y1 = float(corner_a[0]), float(corner_a[1])
|
|
if corner_b is not None:
|
|
x2, y2 = float(corner_b[0]), float(corner_b[1])
|
|
|
|
x1 = float(np.clip(x1, 0.0, 1.0))
|
|
y1 = float(np.clip(y1, 0.0, 1.0))
|
|
x2 = float(np.clip(x2, 0.0, 1.0))
|
|
y2 = float(np.clip(y2, 0.0, 1.0))
|
|
|
|
left = min(x1, x2)
|
|
right = max(x1, x2)
|
|
top = min(y1, y2)
|
|
bottom = max(y1, y2)
|
|
|
|
if square:
|
|
left, top, right, bottom = coerce_physical_square(
|
|
left, top, right, bottom, field.xreal, field.yreal,
|
|
)
|
|
|
|
emit_overlay({
|
|
"kind": "crop_box",
|
|
"section_title": "Preview",
|
|
"image": encode_preview(datafield_to_uint8(field, field.colormap)),
|
|
"x1": left,
|
|
"y1": top,
|
|
"x2": right,
|
|
"y2": bottom,
|
|
"xreal": float(field.xreal),
|
|
"yreal": float(field.yreal),
|
|
"a_locked": corner_a is not None,
|
|
"b_locked": corner_b is not None,
|
|
})
|
|
|
|
px0 = int(np.floor(left * field.xres))
|
|
py0 = int(np.floor(top * field.yres))
|
|
px1 = int(np.ceil(right * field.xres))
|
|
py1 = int(np.ceil(bottom * field.yres))
|
|
|
|
px0 = min(max(px0, 0), field.xres)
|
|
py0 = min(max(py0, 0), field.yres)
|
|
px1 = min(max(px1, px0), field.xres)
|
|
py1 = min(max(py1, py0), field.yres)
|
|
|
|
binary = np.zeros((field.yres, field.xres), dtype=bool)
|
|
binary[py0:py1, px0:px1] = True
|
|
if invert:
|
|
binary = ~binary
|
|
|
|
mask = bool_to_mask(binary)
|
|
|
|
return (mask,)
|