add folder, file nodes and major usability improvements
This commit is contained in:
@@ -7,10 +7,26 @@ before execution begins.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
import json
|
||||
import numpy as np
|
||||
from backend.node_registry import register_node
|
||||
from backend.data_types import (
|
||||
DataField, MeasureTable, COLORMAPS, datafield_to_uint8, image_to_uint8, encode_preview, normalize_for_colormap,
|
||||
DataField,
|
||||
MeasureTable,
|
||||
COLORMAPS,
|
||||
CUSTOM_FILE_FONT,
|
||||
DEFAULT_CUSTOM_COLORMAP_STOPS,
|
||||
SYSTEM_DEFAULT_FONT,
|
||||
colormap_to_uint8,
|
||||
datafield_to_uint8,
|
||||
encode_preview,
|
||||
image_to_uint8,
|
||||
list_overlay_font_choices,
|
||||
normalize_colormap_spec,
|
||||
normalize_font_spec,
|
||||
normalize_for_colormap,
|
||||
render_datafield_preview,
|
||||
resolve_colormap_input,
|
||||
)
|
||||
|
||||
|
||||
@@ -59,15 +75,473 @@ def _scalar_payload(value: float, unit: str = "") -> dict:
|
||||
return payload
|
||||
|
||||
|
||||
_SI_PREFIXES = [
|
||||
(1e24, "Y"),
|
||||
(1e21, "Z"),
|
||||
(1e18, "E"),
|
||||
(1e15, "P"),
|
||||
(1e12, "T"),
|
||||
(1e9, "G"),
|
||||
(1e6, "M"),
|
||||
(1e3, "k"),
|
||||
(1.0, ""),
|
||||
(1e-3, "m"),
|
||||
(1e-6, "u"),
|
||||
(1e-9, "n"),
|
||||
(1e-12, "p"),
|
||||
(1e-15, "f"),
|
||||
(1e-18, "a"),
|
||||
(1e-21, "z"),
|
||||
(1e-24, "y"),
|
||||
]
|
||||
_PREFIXABLE_UNITS = {"m", "s", "A", "V", "W", "Hz", "F", "C", "J", "N", "Pa", "T", "H", "S", "g", "K", "Ohm", "ohm", "Ω"}
|
||||
|
||||
|
||||
def _format_numeric(value: float) -> str:
|
||||
if not np.isfinite(value):
|
||||
return str(value)
|
||||
abs_value = abs(value)
|
||||
if abs_value == 0:
|
||||
return "0"
|
||||
if abs_value >= 1e4 or abs_value < 1e-3:
|
||||
return f"{value:.3e}"
|
||||
return f"{value:.4g}"
|
||||
|
||||
|
||||
def _format_with_unit(value: float, unit: str) -> str:
|
||||
unit = (unit or "").strip()
|
||||
if not unit:
|
||||
return _format_numeric(value)
|
||||
if unit in _PREFIXABLE_UNITS and np.isfinite(value) and value != 0:
|
||||
abs_value = abs(value)
|
||||
for scale, prefix in _SI_PREFIXES:
|
||||
scaled = abs_value / scale
|
||||
if 1 <= scaled < 1000:
|
||||
signed = value / scale
|
||||
return f"{_format_numeric(signed)} {prefix}{unit}"
|
||||
return f"{_format_numeric(value)} {unit}"
|
||||
|
||||
|
||||
def _nice_length(target: float) -> float:
|
||||
if not np.isfinite(target) or target <= 0:
|
||||
return 0.0
|
||||
exponent = np.floor(np.log10(target))
|
||||
base = 10.0 ** exponent
|
||||
for step in (5.0, 2.0, 1.0):
|
||||
candidate = step * base
|
||||
if candidate <= target:
|
||||
return candidate
|
||||
return base
|
||||
|
||||
|
||||
def _display_value_range(field: DataField) -> tuple[float, float, float]:
|
||||
data = np.asarray(field.data, dtype=np.float64)
|
||||
dmin = float(data.min())
|
||||
dmax = float(data.max())
|
||||
if not np.isfinite(dmin) or not np.isfinite(dmax) or dmax <= dmin:
|
||||
return dmin, dmin, dmin
|
||||
|
||||
offset = float(field.display_offset)
|
||||
scale = float(field.display_scale)
|
||||
if not np.isfinite(offset):
|
||||
offset = 0.0
|
||||
if not np.isfinite(scale) or scale <= 0.0:
|
||||
scale = 1.0
|
||||
|
||||
low_norm = float(np.clip(offset, 0.0, 1.0))
|
||||
high_norm = float(np.clip(offset + scale, 0.0, 1.0))
|
||||
if high_norm < low_norm:
|
||||
low_norm, high_norm = high_norm, low_norm
|
||||
mid_norm = 0.5 * (low_norm + high_norm)
|
||||
|
||||
span = dmax - dmin
|
||||
return (
|
||||
dmin + low_norm * span,
|
||||
dmin + mid_norm * span,
|
||||
dmin + high_norm * span,
|
||||
)
|
||||
|
||||
|
||||
def _render_annotation_text(text: str, size_px: int, color: tuple[int, int, int]):
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
|
||||
size_px = max(8, int(round(size_px)))
|
||||
try:
|
||||
font = ImageFont.truetype("DejaVuSans.ttf", size_px)
|
||||
probe = Image.new("RGBA", (1, 1), (0, 0, 0, 0))
|
||||
probe_draw = ImageDraw.Draw(probe)
|
||||
bbox = probe_draw.textbbox((0, 0), text, font=font)
|
||||
width = max(1, bbox[2] - bbox[0])
|
||||
height = max(1, bbox[3] - bbox[1])
|
||||
text_image = Image.new("RGBA", (width, height), (0, 0, 0, 0))
|
||||
text_draw = ImageDraw.Draw(text_image)
|
||||
text_draw.text((-bbox[0], -bbox[1]), text, font=font, fill=(*color, 255))
|
||||
return text_image
|
||||
except Exception:
|
||||
font = ImageFont.load_default()
|
||||
probe = Image.new("L", (1, 1), 0)
|
||||
probe_draw = ImageDraw.Draw(probe)
|
||||
bbox = probe_draw.textbbox((0, 0), text, font=font)
|
||||
width = max(1, bbox[2] - bbox[0])
|
||||
height = max(1, bbox[3] - bbox[1])
|
||||
mask = Image.new("L", (width, height), 0)
|
||||
mask_draw = ImageDraw.Draw(mask)
|
||||
mask_draw.text((-bbox[0], -bbox[1]), text, font=font, fill=255)
|
||||
|
||||
scale = max(1.0, size_px / max(1, height))
|
||||
scaled_width = max(1, int(round(width * scale)))
|
||||
scaled_height = max(1, int(round(height * scale)))
|
||||
resampling = getattr(Image, "Resampling", Image)
|
||||
scaled_mask = mask.resize((scaled_width, scaled_height), resample=resampling.BILINEAR)
|
||||
|
||||
text_image = Image.new("RGBA", (scaled_width, scaled_height), (*color, 0))
|
||||
text_image.putalpha(scaled_mask)
|
||||
return text_image
|
||||
|
||||
|
||||
def _normalize_markup_color(color: object, default: str = "#ffd54f") -> str:
|
||||
if isinstance(color, str):
|
||||
text = color.strip()
|
||||
if len(text) == 4 and text.startswith("#"):
|
||||
text = "#" + "".join(ch * 2 for ch in text[1:])
|
||||
if len(text) == 7 and text.startswith("#"):
|
||||
try:
|
||||
int(text[1:], 16)
|
||||
return text.lower()
|
||||
except ValueError:
|
||||
pass
|
||||
return default
|
||||
|
||||
|
||||
def _parse_markup_shapes(raw_shapes: str | list | None) -> list[dict[str, object]]:
|
||||
if isinstance(raw_shapes, str):
|
||||
try:
|
||||
raw_shapes = json.loads(raw_shapes or "[]")
|
||||
except json.JSONDecodeError:
|
||||
raw_shapes = []
|
||||
|
||||
if not isinstance(raw_shapes, list):
|
||||
return []
|
||||
|
||||
parsed: list[dict[str, object]] = []
|
||||
for shape in raw_shapes:
|
||||
if not isinstance(shape, dict):
|
||||
continue
|
||||
|
||||
kind = str(shape.get("kind", "")).strip().lower()
|
||||
if kind not in {"line", "rectangle", "circle", "arrow"}:
|
||||
continue
|
||||
|
||||
try:
|
||||
x1 = float(shape.get("x1"))
|
||||
y1 = float(shape.get("y1"))
|
||||
x2 = float(shape.get("x2"))
|
||||
y2 = float(shape.get("y2"))
|
||||
width = int(round(float(shape.get("width", 3))))
|
||||
except (TypeError, ValueError):
|
||||
continue
|
||||
|
||||
coords = [x1, y1, x2, y2]
|
||||
if not all(np.isfinite(value) for value in coords):
|
||||
continue
|
||||
|
||||
parsed.append({
|
||||
"kind": kind,
|
||||
"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)),
|
||||
"width": max(1, min(128, width)),
|
||||
"color": _normalize_markup_color(shape.get("color")),
|
||||
})
|
||||
|
||||
return parsed
|
||||
|
||||
|
||||
def _draw_arrow(draw, start: tuple[float, float], end: tuple[float, float], color: str, width: int):
|
||||
dx = end[0] - start[0]
|
||||
dy = end[1] - start[1]
|
||||
length = float(np.hypot(dx, dy))
|
||||
if length <= 1e-6:
|
||||
radius = max(1.0, width / 2.0)
|
||||
draw.ellipse(
|
||||
(start[0] - radius, start[1] - radius, start[0] + radius, start[1] + radius),
|
||||
fill=color,
|
||||
)
|
||||
return
|
||||
|
||||
ux = dx / length
|
||||
uy = dy / length
|
||||
head_length = max(10.0, width * 4.0)
|
||||
head_width = max(8.0, width * 3.0)
|
||||
shaft_end = (
|
||||
end[0] - ux * head_length,
|
||||
end[1] - uy * head_length,
|
||||
)
|
||||
|
||||
draw.line((start, shaft_end), fill=color, width=width)
|
||||
|
||||
px = -uy
|
||||
py = ux
|
||||
left = (
|
||||
shaft_end[0] + px * head_width / 2.0,
|
||||
shaft_end[1] + py * head_width / 2.0,
|
||||
)
|
||||
right = (
|
||||
shaft_end[0] - px * head_width / 2.0,
|
||||
shaft_end[1] - py * head_width / 2.0,
|
||||
)
|
||||
draw.polygon([end, left, right], fill=color)
|
||||
|
||||
|
||||
def _render_markup_image(image: np.ndarray, shapes: list[dict[str, object]]) -> np.ndarray:
|
||||
from PIL import Image, ImageDraw
|
||||
|
||||
base = image_to_uint8(image)
|
||||
if base.ndim == 2:
|
||||
base = np.repeat(base[:, :, np.newaxis], 3, axis=2)
|
||||
|
||||
canvas = Image.fromarray(base.copy())
|
||||
draw = ImageDraw.Draw(canvas)
|
||||
height, width = base.shape[:2]
|
||||
|
||||
for shape in shapes:
|
||||
x1 = float(shape["x1"]) * width
|
||||
y1 = float(shape["y1"]) * height
|
||||
x2 = float(shape["x2"]) * width
|
||||
y2 = float(shape["y2"]) * height
|
||||
color = str(shape["color"])
|
||||
stroke_width = int(shape["width"])
|
||||
kind = str(shape["kind"])
|
||||
|
||||
if kind == "line":
|
||||
draw.line(((x1, y1), (x2, y2)), fill=color, width=stroke_width)
|
||||
elif kind == "rectangle":
|
||||
draw.rectangle((x1, y1, x2, y2), outline=color, width=stroke_width)
|
||||
elif kind == "circle":
|
||||
draw.ellipse((x1, y1, x2, y2), outline=color, width=stroke_width)
|
||||
elif kind == "arrow":
|
||||
_draw_arrow(draw, (x1, y1), (x2, y2), color, stroke_width)
|
||||
|
||||
return np.asarray(canvas, dtype=np.uint8)
|
||||
|
||||
|
||||
@register_node(display_name="Color Map")
|
||||
class ColorMap:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"mode": (["preset", "custom"], {"default": "preset"}),
|
||||
"preset": (list(COLORMAPS), {
|
||||
"default": "viridis",
|
||||
"show_when_widget_value": {"mode": ["preset"]},
|
||||
}),
|
||||
"stops": ("STRING", {
|
||||
"default": json.dumps(list(DEFAULT_CUSTOM_COLORMAP_STOPS)),
|
||||
"colormap_stops": True,
|
||||
"show_when_widget_value": {"mode": ["custom"]},
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("COLORMAP",)
|
||||
RETURN_NAMES = ("colormap",)
|
||||
FUNCTION = "build"
|
||||
CATEGORY = "display"
|
||||
DESCRIPTION = (
|
||||
"Build a reusable colormap. Choose a preset, or create a custom gradient with min/max colours "
|
||||
"and any number of intermediate stops."
|
||||
)
|
||||
|
||||
def build(self, mode: str, preset: str, stops: str | None = None, stops_json: str | None = None) -> tuple:
|
||||
if mode == "preset":
|
||||
return ({"mode": "preset", "preset": normalize_colormap_spec(preset)},)
|
||||
|
||||
try:
|
||||
raw_stops = stops if stops is not None else stops_json
|
||||
stops_data = json.loads(raw_stops or "[]")
|
||||
except json.JSONDecodeError as exc:
|
||||
raise ValueError("Custom colormap stops must be valid JSON.") from exc
|
||||
|
||||
spec = normalize_colormap_spec({"mode": "custom", "stops": stops_data}, fallback=None)
|
||||
if not (isinstance(spec, dict) and spec.get("mode") == "custom"):
|
||||
raise ValueError("Custom colormap must include at least min and max colours.")
|
||||
return (spec,)
|
||||
|
||||
|
||||
@register_node(display_name="Font")
|
||||
class Font:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"family": ([SYSTEM_DEFAULT_FONT, *list_overlay_font_choices(), CUSTOM_FILE_FONT], {
|
||||
"default": SYSTEM_DEFAULT_FONT,
|
||||
}),
|
||||
"font_file": ("FILE_PICKER", {
|
||||
"default": "",
|
||||
"show_when_widget_value": {"family": [CUSTOM_FILE_FONT]},
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("FONT",)
|
||||
RETURN_NAMES = ("font",)
|
||||
FUNCTION = "build"
|
||||
CATEGORY = "display"
|
||||
DESCRIPTION = (
|
||||
"Build a reusable font spec for annotation overlays. Choose a discovered system font, "
|
||||
"use the default fallback stack, or point to a custom font file."
|
||||
)
|
||||
|
||||
def build(self, family: str, font_file: str = "") -> tuple:
|
||||
if family == SYSTEM_DEFAULT_FONT:
|
||||
return (None,)
|
||||
if family == CUSTOM_FILE_FONT:
|
||||
return (normalize_font_spec({"path": font_file}),)
|
||||
return (normalize_font_spec({"family": family}),)
|
||||
|
||||
|
||||
@register_node(display_name="Annotations")
|
||||
class Annotations:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"field": ("DATA_FIELD",),
|
||||
"colormap": (["auto"] + list(COLORMAPS), {"hide_when_input_connected": "colormap_map"}),
|
||||
"show_scale_bar": ("BOOLEAN", {"default": True}),
|
||||
"show_color_map": ("BOOLEAN", {"default": True}),
|
||||
"text_size": ("FLOAT", {
|
||||
"default": 14.0,
|
||||
"min": 6.0,
|
||||
"max": 96.0,
|
||||
"step": 1.0,
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"colormap_map": ("COLORMAP", {"label": "colormap"}),
|
||||
"font": ("FONT",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("DATA_FIELD",)
|
||||
RETURN_NAMES = ("annotated",)
|
||||
FUNCTION = "render"
|
||||
CATEGORY = "display"
|
||||
DESCRIPTION = (
|
||||
"Attach optional publication-style annotations to a DATA_FIELD without flattening the raw data. "
|
||||
"The preview shows a scale bar and/or side colour legend, while downstream field operations keep the underlying AFM values."
|
||||
)
|
||||
|
||||
def render(
|
||||
self,
|
||||
field: DataField,
|
||||
colormap: str,
|
||||
show_scale_bar: bool,
|
||||
show_color_map: bool,
|
||||
text_size: float = 1.0,
|
||||
colormap_map=None,
|
||||
font=None,
|
||||
) -> tuple:
|
||||
resolved_colormap = resolve_colormap_input(
|
||||
colormap,
|
||||
colormap_input=colormap_map,
|
||||
inherited=field.colormap,
|
||||
default="gray",
|
||||
)
|
||||
text_size = float(np.clip(text_size, 6.0, 96.0)) if np.isfinite(text_size) else 14.0
|
||||
out = field.replace(
|
||||
colormap=resolved_colormap,
|
||||
overlays=[
|
||||
*field.overlays,
|
||||
{
|
||||
"kind": "annotation",
|
||||
"show_scale_bar": bool(show_scale_bar),
|
||||
"show_color_map": bool(show_color_map),
|
||||
"text_size": text_size,
|
||||
"font": normalize_font_spec(font),
|
||||
},
|
||||
],
|
||||
)
|
||||
return (out,)
|
||||
|
||||
|
||||
@register_node(display_name="Markup")
|
||||
class Markup:
|
||||
_CUSTOM_PREVIEW = True
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"field": ("DATA_FIELD",),
|
||||
"shape": (["line", "rectangle", "circle", "arrow"], {"default": "line"}),
|
||||
"stroke_color": ("STRING", {"default": "#ffd54f", "color_picker": True}),
|
||||
"stroke_width": ("INT", {"default": 3, "min": 1, "max": 64, "step": 1}),
|
||||
"clear_shapes": ("BUTTON", {"label": "Clear Shapes", "set_widgets": {"markup_shapes": "[]"}}),
|
||||
"markup_shapes": ("STRING", {"default": "[]", "hidden": True}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("DATA_FIELD",)
|
||||
RETURN_NAMES = ("annotated",)
|
||||
FUNCTION = "process"
|
||||
CATEGORY = "display"
|
||||
DESCRIPTION = (
|
||||
"Draw simple vector markup over a DATA_FIELD without flattening the underlying data. "
|
||||
"Choose a shape mode, colour, and stroke width, then drag directly on the preview to place lines, rectangles, circles, or arrows."
|
||||
)
|
||||
|
||||
_broadcast_overlay_fn = None
|
||||
_current_node_id: str = ""
|
||||
|
||||
def process(
|
||||
self,
|
||||
field: DataField,
|
||||
shape: str,
|
||||
stroke_color: str,
|
||||
stroke_width: int,
|
||||
markup_shapes: str,
|
||||
) -> tuple:
|
||||
shapes = _parse_markup_shapes(markup_shapes)
|
||||
out = field.replace(
|
||||
overlays=[
|
||||
*field.overlays,
|
||||
{
|
||||
"kind": "markup",
|
||||
"shapes": shapes,
|
||||
},
|
||||
],
|
||||
)
|
||||
|
||||
if Markup._broadcast_overlay_fn is not None:
|
||||
Markup._broadcast_overlay_fn(
|
||||
Markup._current_node_id,
|
||||
{
|
||||
"kind": "markup",
|
||||
"section_title": "Markup",
|
||||
"image": encode_preview(datafield_to_uint8(field, field.colormap)),
|
||||
"shape": str(shape),
|
||||
"stroke_color": _normalize_markup_color(stroke_color),
|
||||
"stroke_width": max(1, int(stroke_width)),
|
||||
},
|
||||
)
|
||||
|
||||
return (out,)
|
||||
|
||||
|
||||
@register_node(display_name="Preview")
|
||||
class PreviewImage:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"colormap": (["auto"] + list(COLORMAPS),),
|
||||
"colormap": (["auto"] + list(COLORMAPS), {"hide_when_input_connected": "colormap_map"}),
|
||||
},
|
||||
"optional": {
|
||||
"colormap_map": ("COLORMAP", {"label": "colormap"}),
|
||||
"image": ("IMAGE",),
|
||||
"field": ("DATA_FIELD",),
|
||||
}
|
||||
@@ -82,30 +556,35 @@ class PreviewImage:
|
||||
_broadcast_fn = None
|
||||
_current_node_id: str = ""
|
||||
|
||||
def preview(self, colormap: str, image: np.ndarray | None = None, field=None) -> tuple:
|
||||
# Resolve "auto" — use field's colormap if available, else fall back to gray
|
||||
if colormap == "auto":
|
||||
colormap = field.colormap if field is not None else "gray"
|
||||
def preview(
|
||||
self,
|
||||
colormap: str,
|
||||
image: np.ndarray | None = None,
|
||||
field=None,
|
||||
colormap_map=None,
|
||||
) -> tuple:
|
||||
resolved_colormap = resolve_colormap_input(
|
||||
colormap,
|
||||
colormap_input=colormap_map,
|
||||
inherited=field.colormap if field is not None else None,
|
||||
default="gray",
|
||||
)
|
||||
|
||||
# Prefer field if both are connected; accept whichever is provided
|
||||
if field is not None:
|
||||
arr_u8 = datafield_to_uint8(field, colormap)
|
||||
arr_u8 = render_datafield_preview(field, resolved_colormap)
|
||||
elif image is not None:
|
||||
if image.dtype != np.uint8:
|
||||
imin, imax = image.min(), image.max()
|
||||
if imax > imin:
|
||||
norm = (image - imin) / (imax - imin)
|
||||
arr_u8 = image_to_uint8(image)
|
||||
if arr_u8.ndim == 2:
|
||||
if image.dtype == np.uint8:
|
||||
normalized = arr_u8.astype(np.float64) / 255.0
|
||||
else:
|
||||
norm = np.zeros_like(image)
|
||||
arr_u8 = (norm * 255).astype(np.uint8)
|
||||
else:
|
||||
arr_u8 = image
|
||||
|
||||
if arr_u8.ndim == 2 and colormap != "gray":
|
||||
import matplotlib.cm as cm
|
||||
cmap = cm.get_cmap(colormap)
|
||||
rgba = cmap(arr_u8.astype(np.float32) / 255.0)
|
||||
arr_u8 = (rgba[:, :, :3] * 255).astype(np.uint8)
|
||||
imin, imax = image.min(), image.max()
|
||||
if imax > imin:
|
||||
normalized = (image - imin) / (imax - imin)
|
||||
else:
|
||||
normalized = np.zeros_like(image, dtype=np.float64)
|
||||
arr_u8 = colormap_to_uint8(normalized, resolved_colormap)
|
||||
else:
|
||||
raise ValueError("Connect either an IMAGE or DATA_FIELD input to Preview.")
|
||||
|
||||
@@ -124,10 +603,13 @@ class View3D:
|
||||
return {
|
||||
"required": {
|
||||
"field": ("DATA_FIELD",),
|
||||
"colormap": (["auto"] + list(COLORMAPS),),
|
||||
"colormap": (["auto"] + list(COLORMAPS), {"hide_when_input_connected": "colormap_map"}),
|
||||
"z_scale": ("FLOAT", {"default": 1, "min": 0.1, "max": 10.0, "step": 0.05}),
|
||||
"resolution": ("INT", {"default": 128, "min": 32, "max": 512, "step": 16}),
|
||||
}
|
||||
},
|
||||
"optional": {
|
||||
"colormap_map": ("COLORMAP", {"label": "colormap"}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ()
|
||||
@@ -144,9 +626,8 @@ class View3D:
|
||||
|
||||
def render(
|
||||
self, field: DataField,
|
||||
colormap: str, z_scale: float, resolution: int,
|
||||
colormap: str, z_scale: float, resolution: int, colormap_map=None,
|
||||
) -> tuple:
|
||||
import matplotlib.cm as cm
|
||||
import base64
|
||||
|
||||
data = field.data
|
||||
@@ -168,10 +649,13 @@ class View3D:
|
||||
data_max=float(field.data.max()),
|
||||
)
|
||||
|
||||
cmap_name = field.colormap if colormap == "auto" else colormap
|
||||
cmap = cm.get_cmap(cmap_name)
|
||||
rgba = cmap(z_norm) # (ny, nx, 4) float [0,1]
|
||||
colors_u8 = (rgba[:, :, :3] * 255).astype(np.uint8)
|
||||
resolved_colormap = resolve_colormap_input(
|
||||
colormap,
|
||||
colormap_input=colormap_map,
|
||||
inherited=field.colormap,
|
||||
default="gray",
|
||||
)
|
||||
colors_u8 = colormap_to_uint8(z_norm, resolved_colormap)
|
||||
|
||||
# Base64-encode arrays for efficient WS transport
|
||||
z_b64 = base64.b64encode(z.tobytes()).decode()
|
||||
|
||||
@@ -4,11 +4,12 @@ I/O nodes: load and save images and SPM data.
|
||||
|
||||
from __future__ import annotations
|
||||
import os
|
||||
import re
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
|
||||
from backend.node_registry import register_node
|
||||
from backend.data_types import DataField, COLORMAPS, encode_preview, image_to_uint8
|
||||
from backend.data_types import COLORMAPS, DataField, encode_preview, image_to_uint8, resolve_colormap_input
|
||||
from backend.runtime_paths import demo_dir, input_dir, output_dir
|
||||
|
||||
# Resolved at server startup so nodes know where to look
|
||||
@@ -22,6 +23,7 @@ _DEMO_EXTENSIONS = {".png", ".jpg", ".jpeg", ".tiff", ".tif", ".npy", ".npz",
|
||||
_SPM_EXTENSIONS = {".gwy", ".sxm", ".ibw"}
|
||||
_IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".tiff", ".tif", ".bmp"}
|
||||
_ARRAY_EXTENSIONS = {".npy", ".npz"}
|
||||
_PATH_COMPATIBLE_EXTENSIONS = _IMAGE_EXTENSIONS | _ARRAY_EXTENSIONS | _SPM_EXTENSIONS
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -105,6 +107,23 @@ def list_channels(filepath: str) -> list[dict]:
|
||||
return [{"name": "field", "type": "DATA_FIELD"}]
|
||||
|
||||
|
||||
def list_folder_paths(folderpath: str) -> list[dict]:
|
||||
"""Return a folder DIRECTORY plus compatible image/array/SPM FILE_PATH outputs."""
|
||||
path = _resolve_path(folderpath)
|
||||
if not path.exists() or not path.is_dir():
|
||||
return []
|
||||
|
||||
resolved_dir = str(path.resolve())
|
||||
results = [{"name": "directory", "type": "DIRECTORY", "path": resolved_dir}]
|
||||
for entry in sorted(path.iterdir(), key=lambda p: p.name.lower()):
|
||||
if not entry.is_file() or entry.name.startswith("."):
|
||||
continue
|
||||
if entry.suffix.lower() not in _PATH_COMPATIBLE_EXTENSIONS:
|
||||
continue
|
||||
results.append({"name": entry.name, "type": "FILE_PATH", "path": str(entry.resolve())})
|
||||
return results
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# LoadFile (unified loader — replaces LoadImage + LoadSPM)
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -115,9 +134,13 @@ class LoadFile:
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"filename": ("FILE_PICKER", {"default": ""}),
|
||||
"colormap": (list(COLORMAPS),),
|
||||
}
|
||||
"filename": ("FILE_PICKER", {"default": "", "hide_when_input_connected": "path"}),
|
||||
"colormap": (list(COLORMAPS), {"hide_when_input_connected": "colormap_map"}),
|
||||
},
|
||||
"optional": {
|
||||
"colormap_map": ("COLORMAP", {"label": "colormap"}),
|
||||
"path": ("FILE_PATH", {"label": "path"}),
|
||||
},
|
||||
}
|
||||
|
||||
# Default outputs — overridden dynamically by the frontend for multi-channel files
|
||||
@@ -136,26 +159,28 @@ class LoadFile:
|
||||
_broadcast_warning_fn = None
|
||||
_current_node_id = None
|
||||
|
||||
def load(self, filename: str, colormap: str = "viridis"):
|
||||
if not filename or not filename.strip():
|
||||
def load(self, filename: str = "", colormap: str = "viridis", colormap_map=None, path: str | None = None):
|
||||
selected_path = str(path).strip() if path is not None else str(filename).strip()
|
||||
if not selected_path:
|
||||
raise ValueError("No file selected — use Browse to pick a file.")
|
||||
path = _resolve_path(filename)
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"File not found: {path}")
|
||||
if path.is_dir():
|
||||
raise IsADirectoryError(f"Expected a file, got a directory: {path}")
|
||||
path_obj = _resolve_path(selected_path)
|
||||
if not path_obj.exists():
|
||||
raise FileNotFoundError(f"File not found: {path_obj}")
|
||||
if path_obj.is_dir():
|
||||
raise IsADirectoryError(f"Expected a file, got a directory: {path_obj}")
|
||||
|
||||
ext = path.suffix.lower()
|
||||
ext = path_obj.suffix.lower()
|
||||
resolved_colormap = resolve_colormap_input(colormap, colormap_input=colormap_map, default="viridis")
|
||||
|
||||
if ext in _SPM_EXTENSIONS:
|
||||
fields = self._load_spm_all(path, ext)
|
||||
fields = self._load_spm_all(path_obj, ext)
|
||||
for f in fields:
|
||||
f.colormap = colormap
|
||||
f.colormap = resolved_colormap
|
||||
return tuple(fields)
|
||||
|
||||
# Image or array — uncalibrated, single output
|
||||
field = self._load_image_or_array(path, ext)
|
||||
field.colormap = colormap
|
||||
field = self._load_image_or_array(path_obj, ext)
|
||||
field.colormap = resolved_colormap
|
||||
self._send_warning("Uncalibrated data — no physical dimensions.")
|
||||
return (field,)
|
||||
|
||||
@@ -349,8 +374,11 @@ class LoadDemo:
|
||||
return {
|
||||
"required": {
|
||||
"name": (choices,),
|
||||
"colormap": (list(COLORMAPS),),
|
||||
}
|
||||
"colormap": (list(COLORMAPS), {"hide_when_input_connected": "colormap_map"}),
|
||||
},
|
||||
"optional": {
|
||||
"colormap_map": ("COLORMAP", {"label": "colormap"}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("DATA_FIELD",)
|
||||
@@ -359,13 +387,38 @@ class LoadDemo:
|
||||
CATEGORY = "io"
|
||||
DESCRIPTION = "Load a bundled demo file so you can try the app without providing your own data."
|
||||
|
||||
def load(self, name: str, colormap: str = "viridis"):
|
||||
path = DEMO_DIR / name
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"Demo file not found: {name}")
|
||||
|
||||
def load(self, name: str = "", colormap: str = "viridis", colormap_map=None):
|
||||
loader = LoadFile()
|
||||
return loader.load(filename=str(path), colormap=colormap)
|
||||
demo_path = DEMO_DIR / name
|
||||
if not demo_path.exists():
|
||||
raise FileNotFoundError(f"Demo file not found: {name}")
|
||||
return loader.load(filename=str(demo_path), colormap=colormap, colormap_map=colormap_map)
|
||||
|
||||
|
||||
@register_node(display_name="Folder")
|
||||
class Folder:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"folder": ("FOLDER_PICKER", {"default": "", "placement": "top"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("DIRECTORY",)
|
||||
RETURN_NAMES = ("directory",)
|
||||
FUNCTION = "list_files"
|
||||
CATEGORY = "io"
|
||||
DESCRIPTION = (
|
||||
"Pick a folder and output its directory path plus one file socket per compatible image, array, or SPM file inside it. "
|
||||
"Supported files include common images, .npy/.npz arrays, and .gwy/.sxm/.ibw scans."
|
||||
)
|
||||
|
||||
def list_files(self, folder: str) -> tuple:
|
||||
entries = list_folder_paths(folder)
|
||||
if not entries:
|
||||
return tuple()
|
||||
return tuple(item["path"] for item in entries)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -395,6 +448,36 @@ class Coordinate:
|
||||
return ((float(x), float(y)),)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Number
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@register_node(display_name="Number")
|
||||
class Number:
|
||||
"""Provide a fixed scalar value that can feed FLOAT or INT widget sockets."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"value": ("FLOAT", {"default": 0.0, "step": 0.01}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("FLOAT",)
|
||||
RETURN_NAMES = ("value",)
|
||||
FUNCTION = "process"
|
||||
CATEGORY = "io"
|
||||
DESCRIPTION = (
|
||||
"Output a fixed numeric value. "
|
||||
"When connected to FLOAT inputs the exact value is used; "
|
||||
"INT inputs round to the nearest integer at execution time."
|
||||
)
|
||||
|
||||
def process(self, value: float) -> tuple:
|
||||
return (float(value),)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# RangeSlider
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -445,12 +528,32 @@ _MAX_SAVE_FIELDS = 8
|
||||
class SaveImage:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
optional = {}
|
||||
optional = {
|
||||
"directory": ("DIRECTORY", {"label": "directory"}),
|
||||
}
|
||||
for i in range(_MAX_SAVE_FIELDS):
|
||||
optional[f"field_{i}"] = ("DATA_FIELD",)
|
||||
optional[f"field_{i}"] = ("SAVE_LAYER", {"label": f"layer {i + 1}"})
|
||||
optional[f"layer_name_{i}"] = ("STRING", {
|
||||
"default": "",
|
||||
"placeholder": "name",
|
||||
"show_when_input_visible": f"field_{i}",
|
||||
"inline_with_input": f"field_{i}",
|
||||
"hide_label": True,
|
||||
})
|
||||
return {
|
||||
"required": {
|
||||
"filename": ("FILE_PICKER", {"default": ""}),
|
||||
"filename": ("STRING", {
|
||||
"default": "",
|
||||
"placeholder": "filename",
|
||||
"placement": "top",
|
||||
}),
|
||||
"directory_path": ("FOLDER_PICKER", {
|
||||
"default": "",
|
||||
"label": "directory",
|
||||
"placement": "top",
|
||||
"hide_when_input_connected": "directory",
|
||||
"top_socket_input": "directory",
|
||||
}),
|
||||
"format": (["TIFF", "NPZ"],),
|
||||
},
|
||||
"optional": optional,
|
||||
@@ -462,59 +565,130 @@ class SaveImage:
|
||||
OUTPUT_NODE = True
|
||||
MANUAL_TRIGGER = True
|
||||
DESCRIPTION = (
|
||||
"Save one or more DATA_FIELD layers to a single file. "
|
||||
"Connect fields to the inputs — a new slot appears as each is filled. "
|
||||
"TIFF writes float32 multi-page; NPZ writes float64 named arrays. "
|
||||
"Save one or more layers to a single file. "
|
||||
"Each layer input accepts either a DATA_FIELD or an IMAGE, including annotated images. "
|
||||
"Optionally drive the output directory from a folder/path node, while keeping the filename widget for the file name. "
|
||||
"A new slot appears as each one is filled, with a matching per-layer name field. "
|
||||
"TIFF writes multi-page data and stores layer names as page descriptions; "
|
||||
"NPZ writes named arrays using those layer names as keys. "
|
||||
"Click Save to write (does not auto-run)."
|
||||
)
|
||||
|
||||
_broadcast_warning_fn = None
|
||||
_current_node_id = None
|
||||
|
||||
def save(self, filename: str, format: str = "TIFF", **kwargs):
|
||||
# Collect connected fields in order
|
||||
fields = []
|
||||
def save(
|
||||
self,
|
||||
filename: str,
|
||||
directory_path: str = "",
|
||||
format: str = "TIFF",
|
||||
directory: str | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
layers = []
|
||||
layer_names = []
|
||||
for i in range(_MAX_SAVE_FIELDS):
|
||||
f = kwargs.get(f"field_{i}")
|
||||
if f is not None:
|
||||
fields.append(f)
|
||||
layer = kwargs.get(f"field_{i}")
|
||||
if layer is not None:
|
||||
layers.append(layer)
|
||||
layer_names.append(self._resolve_layer_name(kwargs.get(f"layer_name_{i}"), i))
|
||||
|
||||
if not fields:
|
||||
raise ValueError("No fields connected — connect at least one DATA_FIELD input.")
|
||||
if not layers:
|
||||
raise ValueError("No layers connected — connect at least one DATA_FIELD or IMAGE input.")
|
||||
|
||||
if not filename or not filename.strip():
|
||||
raise ValueError("No output path selected — use Browse to pick a location.")
|
||||
|
||||
path = Path(filename)
|
||||
# Ensure parent directory exists
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Force correct extension
|
||||
ext = ".tiff" if format == "TIFF" else ".npz"
|
||||
if path.suffix.lower() != ext:
|
||||
path = path.with_suffix(ext)
|
||||
path = self._resolve_save_path(filename, format, directory, directory_path)
|
||||
|
||||
if format == "TIFF":
|
||||
self._save_tiff(path, fields)
|
||||
self._save_tiff(path, layers, layer_names)
|
||||
else:
|
||||
self._save_npz(path, fields)
|
||||
self._save_npz(path, layers, layer_names)
|
||||
|
||||
self._send_warning(f"Saved {len(fields)} layer(s) to {path.name}")
|
||||
self._send_warning(f"Saved {len(layers)} layer(s) to {path.name}")
|
||||
return ()
|
||||
|
||||
def _save_tiff(self, path: Path, fields: list[DataField]):
|
||||
from PIL import Image
|
||||
images = []
|
||||
for f in fields:
|
||||
images.append(Image.fromarray(f.data.astype(np.float32)))
|
||||
images[0].save(str(path), save_all=True, append_images=images[1:])
|
||||
def _save_tiff(self, path: Path, layers: list[DataField | np.ndarray], layer_names: list[str]):
|
||||
import tifffile
|
||||
|
||||
def _save_npz(self, path: Path, fields: list[DataField]):
|
||||
with tifffile.TiffWriter(str(path)) as tif:
|
||||
for layer, layer_name in zip(layers, layer_names):
|
||||
tif.write(self._layer_array_for_tiff(layer), description=layer_name)
|
||||
|
||||
def _save_npz(self, path: Path, layers: list[DataField | np.ndarray], layer_names: list[str]):
|
||||
arrays = {}
|
||||
for i, f in enumerate(fields):
|
||||
arrays[f"layer_{i}"] = f.data
|
||||
used_keys = set()
|
||||
for i, (layer, layer_name) in enumerate(zip(layers, layer_names)):
|
||||
arrays[self._unique_npz_key(layer_name, used_keys, i)] = self._layer_array_for_npz(layer)
|
||||
np.savez(str(path), **arrays)
|
||||
|
||||
def _resolve_layer_name(self, raw_name: object, index: int) -> str:
|
||||
text = str(raw_name).strip() if raw_name is not None else ""
|
||||
return text or f"layer_{index}"
|
||||
|
||||
def _resolve_save_path(
|
||||
self,
|
||||
filename: str,
|
||||
format: str,
|
||||
directory: str | None,
|
||||
directory_path: str = "",
|
||||
) -> Path:
|
||||
ext = ".tiff" if format == "TIFF" else ".npz"
|
||||
raw_filename = str(filename).strip() if filename is not None else ""
|
||||
raw_directory = str(directory).strip() if directory is not None else ""
|
||||
if not raw_directory:
|
||||
raw_directory = str(directory_path).strip() if directory_path is not None else ""
|
||||
|
||||
if raw_directory:
|
||||
dir_path = Path(raw_directory).expanduser()
|
||||
if dir_path.exists() and not dir_path.is_dir():
|
||||
raise ValueError("Directory input expects a folder path, not a file path.")
|
||||
if not dir_path.exists():
|
||||
if dir_path.suffix:
|
||||
raise ValueError("Directory input expects a folder path, not a file path.")
|
||||
dir_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
filename_part = Path(raw_filename).name if raw_filename else ""
|
||||
if not filename_part:
|
||||
raise ValueError("No output filename selected — enter a file name when using a directory input.")
|
||||
path = dir_path / filename_part
|
||||
else:
|
||||
if not raw_filename:
|
||||
raise ValueError("No output path selected — use Browse to pick a location.")
|
||||
path = Path(raw_filename).expanduser()
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if path.suffix.lower() != ext:
|
||||
path = path.with_suffix(ext)
|
||||
return path
|
||||
|
||||
def _unique_npz_key(self, raw_name: str, used_keys: set[str], index: int) -> str:
|
||||
key = re.sub(r"[^0-9A-Za-z_]+", "_", str(raw_name).strip()).strip("_")
|
||||
if not key:
|
||||
key = f"layer_{index}"
|
||||
if key[0].isdigit():
|
||||
key = f"layer_{key}"
|
||||
|
||||
candidate = key
|
||||
suffix = 2
|
||||
while candidate in used_keys:
|
||||
candidate = f"{key}_{suffix}"
|
||||
suffix += 1
|
||||
used_keys.add(candidate)
|
||||
return candidate
|
||||
|
||||
def _layer_array_for_tiff(self, layer: DataField | np.ndarray) -> np.ndarray:
|
||||
if isinstance(layer, DataField):
|
||||
return np.asarray(layer.data, dtype=np.float32)
|
||||
if isinstance(layer, np.ndarray):
|
||||
return image_to_uint8(layer)
|
||||
raise ValueError(f"Unsupported save layer type: {type(layer).__name__}")
|
||||
|
||||
def _layer_array_for_npz(self, layer: DataField | np.ndarray) -> np.ndarray:
|
||||
if isinstance(layer, DataField):
|
||||
return np.asarray(layer.data)
|
||||
if isinstance(layer, np.ndarray):
|
||||
return np.asarray(layer)
|
||||
raise ValueError(f"Unsupported save layer type: {type(layer).__name__}")
|
||||
|
||||
def _send_warning(self, message: str):
|
||||
fn = SaveImage._broadcast_warning_fn
|
||||
nid = SaveImage._current_node_id
|
||||
|
||||
@@ -143,6 +143,7 @@ class CropResizeField:
|
||||
yreal=(py1 - py0) * field.dy,
|
||||
xoff=field.xoff + px0 * field.dx,
|
||||
yoff=field.yoff + py0 * field.dy,
|
||||
overlays=[],
|
||||
)
|
||||
|
||||
target_width, target_height = self._resolve_target_shape(
|
||||
@@ -217,6 +218,9 @@ class RotateField:
|
||||
"Optionally expand the canvas to keep the full rotated field while preserving the field center."
|
||||
)
|
||||
|
||||
_broadcast_warning_fn = None
|
||||
_current_node_id: str = ""
|
||||
|
||||
def process(
|
||||
self,
|
||||
field: DataField,
|
||||
@@ -224,6 +228,9 @@ class RotateField:
|
||||
interpolation: str,
|
||||
expand_canvas: bool,
|
||||
) -> tuple:
|
||||
if field.overlays:
|
||||
self._send_warning("Rotate clears annotation/markup overlays!")
|
||||
|
||||
angle = float(angle)
|
||||
order_map = {
|
||||
"nearest": 0,
|
||||
@@ -264,9 +271,16 @@ class RotateField:
|
||||
yreal=new_yreal,
|
||||
xoff=center_x - new_xreal / 2.0,
|
||||
yoff=center_y - new_yreal / 2.0,
|
||||
overlays=[],
|
||||
)
|
||||
return (result,)
|
||||
|
||||
def _send_warning(self, message: str):
|
||||
fn = RotateField._broadcast_warning_fn
|
||||
nid = RotateField._current_node_id
|
||||
if fn and nid:
|
||||
fn(nid, message)
|
||||
|
||||
@staticmethod
|
||||
def _rotated_extents(field: DataField, angle: float, expand_canvas: bool) -> tuple[float, float]:
|
||||
if not expand_canvas:
|
||||
|
||||
Reference in New Issue
Block a user