diff --git a/backend/exporters/datafield.py b/backend/exporters/datafield.py index 23d515d..e580ef4 100644 --- a/backend/exporters/datafield.py +++ b/backend/exporters/datafield.py @@ -1,34 +1,44 @@ """ -Exporter for DATA_FIELD values. +Exporter for DATA_FIELD values (single layer or multi-layer stacks). Format choices: -* **TIFF** — 8-bit RGB colormap preview. *Not* round-trippable. Useful for - figures and sharing; opening it back gives you pixels, not physics. -* **TIFF (data)** — float64 array with tono metadata JSON-embedded in the - TIFF ImageDescription tag. Round-trips via the array_image importer once - that importer learns to read the tag (see tests/node_tests/exporters.py). -* **PNG** — 8-bit RGB colormap preview. Not round-trippable. -* **NPZ** — raw ``data`` array only. Not round-trippable (units are dropped). -* **GWY** — Gwyddion native format via the ``gwyfile`` package. Round-trips - and opens directly in Gwyddion. Recommended for "save and come back later". -* **HDF5** — generic HDF5 with one ``/data`` dataset and physical dimensions - as dataset attrs. Round-trips via our generic ``hdf5`` importer. -* **HDF5 (Ergo)** — Asylum Research / Ergo layout with the dataset at - ``Image/DataSet/Resolution 0/Frame 0//Image`` and a sidecar group - ``Image/DataSetInfo/Global/Channels/<title>/ImageDims`` carrying - ``DimScaling`` / ``DimUnits`` / ``DataUnits``. Round-trips via our - ``ergo_hdf5`` importer and opens in Asylum Ergo / Igor. +* **TIFF** — 8-bit RGB colormap preview. *Not* round-trippable and single-layer + only; connect multiple channels and pick "TIFF (data)" for a stack. +* **TIFF (data)** — float64 pixels with tono metadata JSON-embedded in the + TIFF ImageDescription tag. Round-trips and supports multi-page stacks: one + IFD per layer, the first page's description carries a ``{"tono": {...}, + "layers": [...]}`` document. +* **PNG** — 8-bit RGB colormap preview. Single-layer only. +* **NPZ** — for a single layer, writes a plain ``field=...`` key. For a stack, + each layer gets its own key derived from its display name (identifier-safe, + deduplicated). +* **GWY** — Gwyddion native format via the ``gwyfile`` package. A multi-layer + save writes one channel per layer (``/0/data``, ``/1/data``, …), each with + its own title, producing a true multi-channel .gwy file. +* **HDF5** — generic HDF5 with one ``data`` dataset per layer and physical + dimensions as dataset attrs. Round-trips via our generic ``hdf5`` importer, + which picks up every 2-D numeric dataset. +* **HDF5 (Ergo)** — Asylum Research / Ergo layout, one dataset per layer under + ``Image/DataSet/Resolution 0/Frame 0/<title>/Image`` plus a matching sidecar + group ``Image/DataSetInfo/Global/Channels/<title>/ImageDims``. Round-trips + via our ``ergo_hdf5`` importer and opens in Ergo / Igor. + +Mixed layer stacks (DataField + Image) are supported for TIFF (data) and NPZ +only; the physics-carrying formats (GWY, HDF5, HDF5 Ergo) require every layer +to be a DataField and raise a clear error otherwise. """ from __future__ import annotations import json +import re from pathlib import Path +from typing import Any, Sequence import numpy as np -from backend.data_types import DataField, datafield_to_uint8 +from backend.data_types import DataField, datafield_to_uint8, image_to_uint8 from backend.exporters._base import FormatSpec accepted_types: tuple[str, ...] = ("DATA_FIELD",) @@ -43,173 +53,313 @@ FORMATS: dict[str, FormatSpec] = { "HDF5 (Ergo)": FormatSpec(ext=".h5", round_trip=True, label="HDF5 (Asylum Research / Ergo)"), } +# Formats that only make sense for a single layer. When extra layers are +# connected, the Save node raises before we get here, but we keep the check +# defensive so the protocol is enforced at the exporter boundary too. +_SINGLE_LAYER_ONLY: frozenset[str] = frozenset({"TIFF", "PNG"}) + + +def save( + path: Path, + value: DataField, + format_name: str, + *, + extra_layers: Sequence[Any] | None = None, + layer_names: Sequence[str] | None = None, + **_opts, +) -> None: + extras = list(extra_layers or []) + layers: list[Any] = [value, *extras] + names = _resolve_layer_names(layers, layer_names, default_primary=path.stem or "field") + + if extras and format_name in _SINGLE_LAYER_ONLY: + raise ValueError( + f"{format_name} only supports a single layer. Use 'TIFF (data)', " + f"'NPZ', 'GWY', or an HDF5 format for multi-layer saves." + ) -def save(path: Path, value: DataField, format_name: str, **_opts) -> None: if format_name == "TIFF": _save_tiff_preview(path, value) return if format_name == "TIFF (data)": - _save_tiff_data(path, value) + _save_tiff_data(path, layers, names) return if format_name == "PNG": _save_png_preview(path, value) return if format_name == "NPZ": - _save_npz(path, value) + _save_npz(path, layers, names) return if format_name == "GWY": - _save_gwy(path, value) + _save_gwy(path, _require_all_datafields(layers, "GWY"), names) return if format_name == "HDF5": - _save_hdf5_generic(path, value) + _save_hdf5_generic(path, _require_all_datafields(layers, "HDF5"), names) return if format_name == "HDF5 (Ergo)": - _save_hdf5_ergo(path, value) + _save_hdf5_ergo(path, _require_all_datafields(layers, "HDF5 (Ergo)"), names) return raise ValueError(f"Format {format_name!r} is not supported for DATA_FIELD.") +# --------------------------------------------------------------------------- +# Layer helpers +# --------------------------------------------------------------------------- + + +def _resolve_layer_names( + layers: Sequence[Any], + raw_names: Sequence[str] | None, + *, + default_primary: str, +) -> list[str]: + """Fill in layer names, falling back to defaults for blank/missing entries. + + The primary layer (index 0) defaults to ``default_primary`` (usually the + file stem), and each extra layer defaults to ``layer_N+1`` (1-indexed for + humans: "layer 2", "layer 3", …). + """ + raw_names = list(raw_names or []) + out: list[str] = [] + for i in range(len(layers)): + raw = str(raw_names[i]).strip() if i < len(raw_names) and raw_names[i] is not None else "" + if raw: + out.append(raw) + elif i == 0: + out.append(default_primary) + else: + out.append(f"layer_{i + 1}") + return out + + +def _require_all_datafields(layers: Sequence[Any], format_label: str) -> list[DataField]: + """Return the list cast to DataFields, raising if any layer is not one.""" + out: list[DataField] = [] + for i, layer in enumerate(layers): + if not isinstance(layer, DataField): + raise ValueError( + f"{format_label} only supports DataField layers; layer {i + 1} " + f"is a {type(layer).__name__}. Use TIFF (data) or NPZ for mixed stacks." + ) + out.append(layer) + return out + + +def _safe_identifier(name: str, index: int) -> str: + """Turn a free-form layer name into a safe identifier (used as an NPZ key).""" + key = re.sub(r"[^0-9A-Za-z_]+", "_", str(name).strip()).strip("_") + if not key: + key = f"layer_{index + 1}" + if key[0].isdigit(): + key = f"layer_{key}" + return key + + +def _dedupe_keys(raw_keys: Sequence[str]) -> list[str]: + used: set[str] = set() + result: list[str] = [] + for k in raw_keys: + candidate = k + suffix = 2 + while candidate in used: + candidate = f"{k}_{suffix}" + suffix += 1 + used.add(candidate) + result.append(candidate) + return result + + +def _layer_to_float_array(layer: Any) -> np.ndarray: + """Coerce a layer into a float array for TIFF (data). Images are promoted.""" + if isinstance(layer, DataField): + return np.ascontiguousarray(layer.data, dtype=np.float64) + if isinstance(layer, np.ndarray): + # Images are left as-is so multi-channel RGB pages survive the write. + return np.ascontiguousarray(layer) + raise ValueError(f"Unsupported layer type for TIFF (data): {type(layer).__name__}") + + +def _layer_to_npz_array(layer: Any) -> np.ndarray: + if isinstance(layer, DataField): + return np.asarray(layer.data) + if isinstance(layer, np.ndarray): + return np.asarray(layer) + raise ValueError(f"Unsupported layer type for NPZ: {type(layer).__name__}") + + +def _datafield_meta(field: DataField) -> dict: + """Build the JSON-serializable physics metadata dict for a DataField.""" + return { + "xreal": float(field.xreal), + "yreal": float(field.yreal), + "xoff": float(field.xoff), + "yoff": float(field.yoff), + "si_unit_xy": str(field.si_unit_xy), + "si_unit_z": str(field.si_unit_z), + "domain": str(field.domain), + "colormap": field.colormap if isinstance(field.colormap, str) else "viridis", + } + + +# --------------------------------------------------------------------------- +# Per-format writers +# --------------------------------------------------------------------------- + + def _save_tiff_preview(path: Path, field: DataField) -> None: import tifffile tifffile.imwrite(str(path), datafield_to_uint8(field, field.colormap)) -def _save_tiff_data(path: Path, field: DataField) -> None: - """Write the raw float64 data with tono metadata in the ImageDescription tag. +def _save_tiff_data(path: Path, layers: Sequence[Any], names: Sequence[str]) -> None: + """Write the raw pixels as a multi-page TIFF with tono metadata. - The description is a JSON document of shape ``{"tono": {...}}`` so future - schema extensions can coexist with other tools' TIFF metadata. Only the - fields needed to reconstruct physical coordinates and z-scaling are - embedded; display state (colormap, display_scale) is intentionally out of - scope — this format is for data, not styling. + The ImageDescription tag on the first page carries a JSON document of + shape ``{"tono": {"version": 1, "layers": [{...}, {...}]}}``. Each entry in + ``layers`` gives the per-layer physics (xreal/yreal/xoff/yoff/units/domain) + and its display name so a future multi-layer importer can reconstruct the + whole stack. Non-DataField layers (plain images) get a minimal entry with + just the name and dtype — they're pixels, not physics. """ import tifffile - meta = { - "tono": { - "version": 1, - "xreal": float(field.xreal), - "yreal": float(field.yreal), - "xoff": float(field.xoff), - "yoff": float(field.yoff), - "si_unit_xy": str(field.si_unit_xy), - "si_unit_z": str(field.si_unit_z), - "domain": str(field.domain), - "colormap": field.colormap if isinstance(field.colormap, str) else "viridis", - } - } - tifffile.imwrite( - str(path), - np.ascontiguousarray(field.data, dtype=np.float64), - description=json.dumps(meta, separators=(",", ":")), + per_layer_meta: list[dict] = [] + for layer, layer_name in zip(layers, names): + if isinstance(layer, DataField): + entry = {"name": layer_name, "kind": "data_field", **_datafield_meta(layer)} + else: + arr = np.asarray(layer) + entry = {"name": layer_name, "kind": "image", "dtype": str(arr.dtype), "shape": list(arr.shape)} + per_layer_meta.append(entry) + + description = json.dumps( + {"tono": {"version": 1, "layers": per_layer_meta}}, + separators=(",", ":"), ) + with tifffile.TiffWriter(str(path)) as tif: + for i, (layer, layer_name) in enumerate(zip(layers, names)): + arr = _layer_to_float_array(layer) + # Full metadata document lives on the first page; subsequent pages + # carry only their display name so readers that walk IFDs see + # something meaningful per channel. + page_desc = description if i == 0 else layer_name + tif.write(arr, description=page_desc) + def _save_png_preview(path: Path, field: DataField) -> None: from PIL import Image Image.fromarray(datafield_to_uint8(field, field.colormap)).save(str(path)) -def _save_npz(path: Path, field: DataField) -> None: - np.savez(str(path), field=np.asarray(field.data)) +def _save_npz(path: Path, layers: Sequence[Any], names: Sequence[str]) -> None: + if len(layers) == 1: + # Single-layer: keep the historical `field` key so nothing that reads + # existing tono .npz outputs breaks. + np.savez(str(path), field=_layer_to_npz_array(layers[0])) + return + raw_keys = [_safe_identifier(name, i) for i, name in enumerate(names)] + keys = _dedupe_keys(raw_keys) + arrays = {key: _layer_to_npz_array(layer) for key, layer in zip(keys, layers)} + np.savez(str(path), **arrays) -def _save_gwy(path: Path, field: DataField) -> None: - """Write a single-channel .gwy file via the gwyfile package.""" +def _save_gwy(path: Path, fields: list[DataField], names: Sequence[str]) -> None: + """Write an N-channel .gwy file via the gwyfile package.""" from gwyfile.objects import GwyContainer, GwyDataField, GwySIUnit - # gwyfile's GwyDataField ctor expects the data array and physical extents. - # si_unit_xy / si_unit_z accept a GwySIUnit wrapper with a .unitstr field. - gwy_field = GwyDataField( - np.ascontiguousarray(field.data, dtype=np.float64), - xreal=float(field.xreal), - yreal=float(field.yreal), - xoff=float(field.xoff), - yoff=float(field.yoff), - si_unit_xy=GwySIUnit(unitstr=str(field.si_unit_xy or "")), - si_unit_z=GwySIUnit(unitstr=str(field.si_unit_z or "")), - ) - title = path.stem or "field" - container = GwyContainer({ - "/0/data": gwy_field, - "/0/data/title": title, - }) - container.tofile(str(path)) + container_data: dict[str, Any] = {} + for i, (field, title) in enumerate(zip(fields, names)): + gwy_field = GwyDataField( + np.ascontiguousarray(field.data, dtype=np.float64), + xreal=float(field.xreal), + yreal=float(field.yreal), + xoff=float(field.xoff), + yoff=float(field.yoff), + si_unit_xy=GwySIUnit(unitstr=str(field.si_unit_xy or "")), + si_unit_z=GwySIUnit(unitstr=str(field.si_unit_z or "")), + ) + container_data[f"/{i}/data"] = gwy_field + container_data[f"/{i}/data/title"] = title + GwyContainer(container_data).tofile(str(path)) -def _save_hdf5_generic(path: Path, field: DataField) -> None: - """Write a single dataset ``/data`` with physical dimensions as dataset attrs. +def _save_hdf5_generic(path: Path, fields: list[DataField], names: Sequence[str]) -> None: + """Write one HDF5 dataset per layer with physical dims as dataset attrs. - The layout is the mirror of :mod:`backend.importers.hdf5`: any 2-D numeric - dataset is picked up and its attrs (``xreal``, ``yreal``, ``xoff``, ``yoff``, - ``si_unit_xy``, ``si_unit_z``) reconstruct the DataField. + Single-layer saves use ``/data`` for backward compatibility with the + tests that read the original layout; multi-layer saves use one + top-level dataset per channel, keyed by the safe-identifier form of its + name and deduplicated against collisions. """ import h5py - arr = np.ascontiguousarray(field.data, dtype=np.float64) with h5py.File(str(path), "w") as f: - ds = f.create_dataset("data", data=arr) - ds.attrs["xreal"] = float(field.xreal) - ds.attrs["yreal"] = float(field.yreal) - ds.attrs["xoff"] = float(field.xoff) - ds.attrs["yoff"] = float(field.yoff) - ds.attrs["si_unit_xy"] = str(field.si_unit_xy or "") - ds.attrs["si_unit_z"] = str(field.si_unit_z or "") + if len(fields) == 1: + _write_hdf5_dataset(f, "data", fields[0]) + return + raw_keys = [_safe_identifier(name, i) for i, name in enumerate(names)] + keys = _dedupe_keys(raw_keys) + for key, field in zip(keys, fields): + _write_hdf5_dataset(f, key, field) -def _save_hdf5_ergo(path: Path, field: DataField) -> None: - """Write an Asylum Research / Ergo-compatible HDF5 file. +def _write_hdf5_dataset(h5file: Any, name: str, field: DataField) -> None: + arr = np.ascontiguousarray(field.data, dtype=np.float64) + ds = h5file.create_dataset(name, data=arr) + ds.attrs["xreal"] = float(field.xreal) + ds.attrs["yreal"] = float(field.yreal) + ds.attrs["xoff"] = float(field.xoff) + ds.attrs["yoff"] = float(field.yoff) + ds.attrs["si_unit_xy"] = str(field.si_unit_xy or "") + ds.attrs["si_unit_z"] = str(field.si_unit_z or "") - The layout mirrors :mod:`backend.importers.ergo_hdf5`: - * The image dataset lives at - ``Image/DataSet/Resolution 0/Frame 0/<title>/Image`` — the second-to-last - path component is the channel name that the importer keys off. - * A sidecar group at - ``Image/DataSetInfo/Global/Channels/<title>/ImageDims`` carries - ``DimScaling`` (a (2, 2) array of absolute physical ranges, Y-first), - ``DimUnits`` (``[Y_unit, X_unit]``), and ``DataUnits`` (Z unit string). +def _save_hdf5_ergo(path: Path, fields: list[DataField], names: Sequence[str]) -> None: + """Write an Asylum Research / Ergo-compatible HDF5 file (N channels). - This makes the file openable by Asylum Ergo / Igor and round-trippable - through our ergo_hdf5 importer. + Each channel gets its own dataset at + ``Image/DataSet/Resolution 0/Frame 0/<title>/Image`` with a matching + sidecar group ``Image/DataSetInfo/Global/Channels/<title>/ImageDims`` + carrying ``DimScaling`` / ``DimUnits`` / ``DataUnits``. The channel + names are the dedupe-safe form of each layer name. Opens in Ergo / Igor + and round-trips through :mod:`backend.importers.ergo_hdf5`. """ import h5py - arr = np.ascontiguousarray(field.data, dtype=np.float64) - title = path.stem or "field" - - x_start = float(field.xoff) - x_end = float(field.xoff) + float(field.xreal) - y_start = float(field.yoff) - y_end = float(field.yoff) + float(field.yreal) - # DimScaling is stored Y-first to match the importer's expectations - # (see ergo_hdf5.py:110-113). - dim_scaling = np.array( - [[y_start, y_end], [x_start, x_end]], - dtype=np.float64, - ) - # DimUnits is [Y_unit, X_unit]; the importer takes the X (second) entry - # as the canonical lateral unit (see ergo_hdf5.py:129-135). - xy_unit = str(field.si_unit_xy or "m") - z_unit = str(field.si_unit_z or "") - dim_units = np.array([xy_unit, xy_unit], dtype=h5py.string_dtype()) + raw_keys = [_safe_identifier(name, i) for i, name in enumerate(names)] + titles = _dedupe_keys(raw_keys) with h5py.File(str(path), "w") as f: - ds = f.create_dataset( - f"Image/DataSet/Resolution 0/Frame 0/{title}/Image", - data=arr, - ) - # Also write the generic attrs so non-Ergo readers still see physics. - ds.attrs["xreal"] = float(field.xreal) - ds.attrs["yreal"] = float(field.yreal) - ds.attrs["xoff"] = float(field.xoff) - ds.attrs["yoff"] = float(field.yoff) - ds.attrs["si_unit_xy"] = xy_unit - ds.attrs["si_unit_z"] = z_unit + for field, title in zip(fields, titles): + arr = np.ascontiguousarray(field.data, dtype=np.float64) + ds = f.create_dataset( + f"Image/DataSet/Resolution 0/Frame 0/{title}/Image", + data=arr, + ) + ds.attrs["xreal"] = float(field.xreal) + ds.attrs["yreal"] = float(field.yreal) + ds.attrs["xoff"] = float(field.xoff) + ds.attrs["yoff"] = float(field.yoff) + xy_unit = str(field.si_unit_xy or "m") + z_unit = str(field.si_unit_z or "") + ds.attrs["si_unit_xy"] = xy_unit + ds.attrs["si_unit_z"] = z_unit - dims_grp = f.create_group( - f"Image/DataSetInfo/Global/Channels/{title}/ImageDims" - ) - dims_grp.attrs["DimScaling"] = dim_scaling - dims_grp.attrs["DimUnits"] = dim_units - dims_grp.attrs["DataUnits"] = z_unit + x_start = float(field.xoff) + x_end = float(field.xoff) + float(field.xreal) + y_start = float(field.yoff) + y_end = float(field.yoff) + float(field.yreal) + # DimScaling is Y-first to match the importer (ergo_hdf5.py:110-113). + dim_scaling = np.array( + [[y_start, y_end], [x_start, x_end]], + dtype=np.float64, + ) + dim_units = np.array([xy_unit, xy_unit], dtype=h5py.string_dtype()) + + dims_grp = f.create_group( + f"Image/DataSetInfo/Global/Channels/{title}/ImageDims" + ) + dims_grp.attrs["DimScaling"] = dim_scaling + dims_grp.attrs["DimUnits"] = dim_units + dims_grp.attrs["DataUnits"] = z_unit diff --git a/backend/exporters/image.py b/backend/exporters/image.py index 2557ece..60bc9f5 100644 --- a/backend/exporters/image.py +++ b/backend/exporters/image.py @@ -5,11 +5,17 @@ Images are raw pixel arrays — no physical calibration by design — so none of the formats here round-trip dimensions. PNG/TIFF convert to uint8 via the same image_to_uint8 helper the preview pipeline uses; NPZ preserves the raw array. + +Multi-layer stacks are supported for TIFF (multi-page uint8) and NPZ (one +named array per layer). PNG is single-layer only and raises if extra layers +are connected. """ from __future__ import annotations +import re from pathlib import Path +from typing import Any, Sequence import numpy as np @@ -24,18 +30,100 @@ FORMATS: dict[str, FormatSpec] = { "NPZ": FormatSpec(ext=".npz", round_trip=False, label="NumPy (.npz)"), } +_SINGLE_LAYER_ONLY: frozenset[str] = frozenset({"PNG"}) + + +def save( + path: Path, + value: np.ndarray, + format_name: str, + *, + extra_layers: Sequence[Any] | None = None, + layer_names: Sequence[str] | None = None, + **_opts, +) -> None: + extras = list(extra_layers or []) + layers: list[Any] = [value, *extras] + names = _resolve_layer_names(len(layers), layer_names, default_primary=path.stem or "image") + + if extras and format_name in _SINGLE_LAYER_ONLY: + raise ValueError( + f"{format_name} only supports a single layer. Use 'TIFF' or 'NPZ' " + f"for multi-layer image saves." + ) -def save(path: Path, value: np.ndarray, format_name: str, **_opts) -> None: - arr = np.asarray(value) if format_name == "PNG": from PIL import Image - Image.fromarray(image_to_uint8(arr)).save(str(path)) + Image.fromarray(image_to_uint8(np.asarray(value))).save(str(path)) return if format_name == "TIFF": - import tifffile - tifffile.imwrite(str(path), image_to_uint8(arr)) + _save_tiff(path, layers, names) return if format_name == "NPZ": - np.savez(str(path), image=arr) + _save_npz(path, layers, names) return raise ValueError(f"Format {format_name!r} is not supported for IMAGE.") + + +def _save_tiff(path: Path, layers: Sequence[Any], names: Sequence[str]) -> None: + import tifffile + + if len(layers) == 1: + tifffile.imwrite(str(path), image_to_uint8(np.asarray(layers[0]))) + return + with tifffile.TiffWriter(str(path)) as tif: + for layer, layer_name in zip(layers, names): + tif.write(image_to_uint8(np.asarray(layer)), description=layer_name) + + +def _save_npz(path: Path, layers: Sequence[Any], names: Sequence[str]) -> None: + if len(layers) == 1: + # Preserve the single-layer key used by the legacy test suite. + np.savez(str(path), image=np.asarray(layers[0])) + return + raw_keys = [_safe_identifier(name, i) for i, name in enumerate(names)] + keys = _dedupe_keys(raw_keys) + arrays = {key: np.asarray(layer) for key, layer in zip(keys, layers)} + np.savez(str(path), **arrays) + + +def _resolve_layer_names( + count: int, + raw_names: Sequence[str] | None, + *, + default_primary: str, +) -> list[str]: + raw_names = list(raw_names or []) + out: list[str] = [] + for i in range(count): + raw = str(raw_names[i]).strip() if i < len(raw_names) and raw_names[i] is not None else "" + if raw: + out.append(raw) + elif i == 0: + out.append(default_primary) + else: + out.append(f"layer_{i + 1}") + return out + + +def _safe_identifier(name: str, index: int) -> str: + key = re.sub(r"[^0-9A-Za-z_]+", "_", str(name).strip()).strip("_") + if not key: + key = f"layer_{index + 1}" + if key[0].isdigit(): + key = f"layer_{key}" + return key + + +def _dedupe_keys(raw_keys: Sequence[str]) -> list[str]: + used: set[str] = set() + result: list[str] = [] + for k in raw_keys: + candidate = k + suffix = 2 + while candidate in used: + candidate = f"{k}_{suffix}" + suffix += 1 + used.add(candidate) + result.append(candidate) + return result diff --git a/backend/node_menu.py b/backend/node_menu.py index b32c27c..014703e 100644 --- a/backend/node_menu.py +++ b/backend/node_menu.py @@ -33,7 +33,6 @@ MENU_LAYOUT: dict[str, list[str]] = { "ValueIO", "PrintTable", "Save", - "SaveImage", "Shade", "PresentationOps", ], diff --git a/backend/nodes/save.py b/backend/nodes/save.py index e885976..7db5d52 100644 --- a/backend/nodes/save.py +++ b/backend/nodes/save.py @@ -2,6 +2,7 @@ from __future__ import annotations import tempfile from pathlib import Path +from typing import Any from backend.node_registry import register_node from backend.execution_context import emit_warning, emit_file_download @@ -11,9 +12,15 @@ from backend.exporters import ( resolve_path, type_name_for_value, ) +from backend.nodes.helpers import _MAX_SAVE_FIELDS DOWNLOAD_DIR = Path(tempfile.gettempdir()) / "tono-downloads" +# Source types that expand into a layer stack (i.e., the Save node grows +# extra field_N inputs). Any other type (FLOAT, LINE, MESH, …) is a single +# value; no stacking UI is shown. +_STACKABLE_SOURCE_TYPES: tuple[str, ...] = ("DATA_FIELD", "IMAGE", "ANNOTATION_SOURCE") + def _choices_by_source_type() -> dict[str, list[str]]: """Build the format dropdown's source-type map from the exporter registry. @@ -39,6 +46,43 @@ class Save: @classmethod def INPUT_TYPES(cls): choices = _choices_by_source_type() + + optional: dict[str, Any] = { + "plot_title": ("STRING", { + "default": "", + "placeholder": "plot title (optional)", + "label": "title", + "show_when_source_type": {"value": ["LINE"]}, + }), + # Name widget for the primary (value) layer. Only surfaces once + # the stack grows beyond one layer, so single-value saves stay + # clutter-free. + "primary_name": ("STRING", { + "default": "", + "placeholder": "name", + "show_when_input_visible": "field_0", + "inline_with_input": "value", + "hide_label": True, + }), + } + # Extra layer sockets for stackable source types. The frontend + # progressive-reveal block keys off `field_N` and only shows slot N + # once slot N-1 is connected; we further gate every slot on `value` + # being a stackable source type via `show_when_source_type`. + for i in range(_MAX_SAVE_FIELDS): + optional[f"field_{i}"] = ("DATA_FIELD", { + "label": f"layer {i + 2}", # primary is layer 1 + "accepted_types": ["IMAGE", "ANNOTATION_SOURCE"], + "show_when_source_type": {"value": list(_STACKABLE_SOURCE_TYPES)}, + }) + 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": ("STRING", { @@ -64,14 +108,7 @@ class Save: "source_type_input": "value", }), }, - "optional": { - "plot_title": ("STRING", { - "default": "", - "placeholder": "plot title (optional)", - "label": "title", - "show_when_source_type": {"value": ["LINE"]}, - }), - }, + "optional": optional, } OUTPUTS = () @@ -80,12 +117,18 @@ class Save: OUTPUT_NODE = True MANUAL_TRIGGER = True DESCRIPTION = ( - "Save a single graph value to disk. Supports fields, images, lines, tables, scalars, " - "and 3D meshes. Use 'GWY' or 'TIFF (data)' for DataFields you want to re-open later " - "with their physical units preserved." + "Save one or more graph values to disk. A single value works for every type " + "(fields, images, lines, tables, scalars, meshes). For DataFields and Images, " + "additional layer slots appear as you connect each one, letting you write " + "multi-channel TIFF, NPZ, GWY, or HDF5 stacks from a single node. " + "Use 'GWY' or 'TIFF (data)' when you need to re-open the result with its " + "physical units preserved." ) - KEYWORDS = ("export", "write", "download", "png", "tiff", "csv", "json", "npz", "obj", "stl", "gwy") + KEYWORDS = ( + "export", "write", "download", "png", "tiff", "csv", "json", "npz", + "obj", "stl", "gwy", "hdf5", "layers", "stack", "channels", + ) def save( self, @@ -93,12 +136,62 @@ class Save: format: str, value, plot_title: str = "", + primary_name: str = "", + **kwargs, ): type_name = type_name_for_value(value) module, spec = get_exporter(type_name, format) path = resolve_path(filename, spec, DOWNLOAD_DIR) - module.save(path, value, format, plot_title=plot_title) + + extra_layers, layer_names = self._collect_extra_layers( + type_name, primary_name, kwargs, + ) + + module.save( + path, + value, + format, + plot_title=plot_title, + extra_layers=extra_layers, + layer_names=layer_names, + ) emit_warning(f"Saved to {path.name}") emit_file_download(str(path)) return () + + def _collect_extra_layers( + self, + type_name: str, + primary_name: str, + kwargs: dict[str, Any], + ) -> tuple[list[Any], list[str]]: + """Pull field_N + layer_name_N from kwargs into parallel lists. + + Only applies when the primary value is a stackable source type; for + anything else (LINE, FLOAT, MESH_MODEL, tables) any stray field_N + kwargs are ignored — the frontend hides those sockets in that case + and the backend treats it as a single-value save. + """ + if type_name not in _STACKABLE_SOURCE_TYPES: + return [], [] + + extras: list[Any] = [] + extra_names: list[str] = [] + # Preserve the on-node order: iterate field_0, field_1, …, stopping at + # the first hole. An unconnected slot in the middle would be a UI bug, + # but bailing early keeps the saved stack matching what the user sees. + for i in range(_MAX_SAVE_FIELDS): + layer = kwargs.get(f"field_{i}") + if layer is None: + break + extras.append(layer) + extra_names.append(str(kwargs.get(f"layer_name_{i}", "") or "").strip()) + + if not extras: + return [], [] + + # Full names list starts with the primary's name (empty → exporter + # substitutes path.stem) and then each extra in order. + names = [str(primary_name or "").strip(), *extra_names] + return extras, names diff --git a/backend/nodes/save_layers.py b/backend/nodes/save_layers.py deleted file mode 100644 index e16eb6d..0000000 --- a/backend/nodes/save_layers.py +++ /dev/null @@ -1,185 +0,0 @@ -from __future__ import annotations -import re -import numpy as np -from pathlib import Path - -from backend.node_registry import register_node -from backend.execution_context import emit_warning, emit_file_download -from backend.data_types import DataField, image_to_uint8 -from backend.nodes.helpers import _MAX_SAVE_FIELDS - - -@register_node(display_name="Save Layers") -class SaveImage: - @classmethod - def INPUT_TYPES(cls): - optional = { - "directory": ("DIRECTORY", {"label": "directory"}), - } - for i in range(_MAX_SAVE_FIELDS): - optional[f"field_{i}"] = ("DATA_FIELD", { - "label": f"layer {i + 1}", - "accepted_types": ["IMAGE", "ANNOTATION_SOURCE"], - }) - 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": ("STRING", { - "default": "", - "placeholder": "filename", - "placement": "top", - }), - "directory_path": ("STRING", { - "default": "", - "label": "directory", - "placeholder": "directory (optional, desktop only)", - "placement": "top", - "hide_when_input_connected": "directory", - "top_socket_input": "directory", - }), - "format": (["TIFF", "NPZ"],), - }, - "optional": optional, - } - - OUTPUTS = () - FUNCTION = "save" - - OUTPUT_NODE = True - MANUAL_TRIGGER = True - DESCRIPTION = ( - "Save one or more image/field 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. " - "Use this for composing multi-channel stacks. 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)." - ) - - KEYWORDS = ("export", "write", "multipage", "stack", "tiff", "npz", "channels") - - 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): - 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 layers: - raise ValueError("No layers connected — connect at least one DATA_FIELD or IMAGE input.") - - path = self._resolve_save_path(filename, format, directory, directory_path) - - if format == "TIFF": - self._save_tiff(path, layers, layer_names) - else: - self._save_npz(path, layers, layer_names) - - emit_warning(f"Saved {len(layers)} layer(s) to {path.name}") - emit_file_download(str(path)) - return () - - def _save_tiff(self, path: Path, layers: list[DataField | np.ndarray], layer_names: list[str]): - import tifffile - - 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 = {} - 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 filename selected — enter a file name.") - candidate = Path(raw_filename).expanduser() - if candidate.is_absolute(): - candidate.parent.mkdir(parents=True, exist_ok=True) - path = candidate - else: - from backend.nodes.save import DOWNLOAD_DIR - DOWNLOAD_DIR.mkdir(parents=True, exist_ok=True) - path = DOWNLOAD_DIR / candidate.name - - 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__}") diff --git a/frontend/src/CustomNode.tsx b/frontend/src/CustomNode.tsx index 19bc57f..e714b8a 100644 --- a/frontend/src/CustomNode.tsx +++ b/frontend/src/CustomNode.tsx @@ -1110,11 +1110,24 @@ function CustomNode({ id, data }: { id: string; data: NodeData }) { for (const [name, spec] of Object.entries(optional)) { const [type, opts] = getSpecTypeAndOptions(spec as InputSpec); if (isProgressive && isDataSocketSpec(spec as InputSpec)) { - // Progressive: show this slot only if it's the first or the previous is connected + // Progressive: show this slot only if it's the first or the previous + // is connected. If the socket also carries `show_when_source_type`, + // the gating input must be connected to a matching source type before + // any slot in the chain is revealed — this lets Save's layer stack + // stay hidden until `value` is a DataField or Image. const match = name.match(/^field_(\d+)$/); if (match) { const idx = parseInt(match[1], 10); - if (idx === 0 || (connectedInputs && connectedInputs.has(`field_${idx - 1}`))) { + const sourceTypeRules = opts?.show_when_source_type; + let sourceTypeOk = true; + if (sourceTypeRules && typeof sourceTypeRules === 'object') { + const gateInput = Object.keys(sourceTypeRules)[0]; + const allowed = Array.isArray(sourceTypeRules[gateInput]) ? sourceTypeRules[gateInput] : []; + const actualSourceType = connectedSourceTypes?.[gateInput] ?? null; + sourceTypeOk = actualSourceType != null && allowed.includes(actualSourceType); + } + const progressiveOk = idx === 0 || (connectedInputs && connectedInputs.has(`field_${idx - 1}`)); + if (sourceTypeOk && progressiveOk) { dataInputs.push({ name, type, label: formatUiLabel(opts?.label || name) }); visibleInputNames.add(name); } diff --git a/tests/node_tests/exporters.py b/tests/node_tests/exporters.py index 1f56704..ff99869 100644 --- a/tests/node_tests/exporters.py +++ b/tests/node_tests/exporters.py @@ -188,12 +188,18 @@ def test_datafield_tiff_data_round_trip(): assert arr.shape == field.data.shape assert np.allclose(arr, field.data) + # Per-layer metadata lives under tono.layers[*]; a single-layer save + # still produces the same shape, just with one entry. meta = json.loads(desc)["tono"] - assert meta["xreal"] == field.xreal - assert meta["yreal"] == field.yreal - assert meta["si_unit_xy"] == "m" - assert meta["si_unit_z"] == "V" - assert meta["domain"] == "spatial" + assert meta["version"] == 1 + assert len(meta["layers"]) == 1 + layer0 = meta["layers"][0] + assert layer0["kind"] == "data_field" + assert layer0["xreal"] == field.xreal + assert layer0["yreal"] == field.yreal + assert layer0["si_unit_xy"] == "m" + assert layer0["si_unit_z"] == "V" + assert layer0["domain"] == "spatial" def test_datafield_hdf5_generic_round_trip(): @@ -282,6 +288,268 @@ def test_datafield_hdf5_ergo_round_trip(): assert rf.si_unit_z == "N" +def test_save_multi_layer_tiff_data(): + """TIFF (data) with extra layers writes multi-page float64 with per-layer metadata.""" + import tifffile + + from backend.nodes.save import Save + + rng = np.random.default_rng(41) + primary = DataField( + data=rng.standard_normal((16, 20)).astype(np.float64) * 1e-9, + xreal=3e-6, yreal=2e-6, si_unit_xy="m", si_unit_z="m", + ) + layer2 = DataField( + data=rng.standard_normal((16, 20)).astype(np.float64) * 1e-12, + xreal=3e-6, yreal=2e-6, si_unit_xy="m", si_unit_z="N", + ) + layer3 = DataField( + data=rng.standard_normal((16, 20)).astype(np.float64), + xreal=3e-6, yreal=2e-6, si_unit_xy="m", si_unit_z="V", + ) + + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "stack" + Save().save( + filename=str(path), + format="TIFF (data)", + value=primary, + field_0=layer2, + field_1=layer3, + primary_name="height", + layer_name_0="force", + layer_name_1="potential", + ) + out_path = path.with_suffix(".tiff") + assert out_path.exists() + + with tifffile.TiffFile(out_path) as tif: + assert len(tif.pages) == 3 + meta = json.loads(tif.pages[0].tags["ImageDescription"].value)["tono"] + assert len(meta["layers"]) == 3 + assert [layer["name"] for layer in meta["layers"]] == ["height", "force", "potential"] + assert meta["layers"][1]["si_unit_z"] == "N" + assert meta["layers"][2]["si_unit_z"] == "V" + assert tif.pages[0].asarray().shape == (16, 20) + assert tif.pages[1].asarray().shape == (16, 20) + assert np.allclose(tif.pages[0].asarray(), primary.data) + assert np.allclose(tif.pages[2].asarray(), layer3.data) + + +def test_save_multi_layer_npz_named_keys(): + """Multi-layer NPZ uses safe-identifier keys from layer names.""" + from backend.nodes.save import Save + + rng = np.random.default_rng(47) + primary = DataField(data=rng.standard_normal((8, 8)).astype(np.float64)) + layer2 = DataField(data=rng.standard_normal((8, 8)).astype(np.float64)) + annotated = np.zeros((12, 12, 3), dtype=np.uint8) + annotated[..., 0] = 255 + + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "stack" + Save().save( + filename=str(path), + format="NPZ", + value=primary, + field_0=layer2, + field_1=annotated, + primary_name="height map", + layer_name_0="force-retrace", + layer_name_1="annotated overview", + ) + out_path = path.with_suffix(".npz") + assert out_path.exists() + + npz = np.load(out_path) + # Non-identifier characters collapse to underscores. + assert set(npz.files) == {"height_map", "force_retrace", "annotated_overview"} + assert np.allclose(npz["height_map"], primary.data) + assert np.allclose(npz["force_retrace"], layer2.data) + assert np.array_equal(npz["annotated_overview"], annotated) + + +def test_save_multi_layer_tiff_preview_rejected(): + """Single-layer-only formats must reject extra layers with a clear error.""" + from backend.nodes.save import Save + + field_a = DataField(data=np.zeros((4, 4))) + field_b = DataField(data=np.ones((4, 4))) + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "preview" + try: + Save().save( + filename=str(path), + format="TIFF", # preview format, single-layer only + value=field_a, + field_0=field_b, + ) + assert False, "TIFF preview must reject extra layers" + except ValueError as exc: + assert "single layer" in str(exc).lower() + + try: + Save().save( + filename=str(path), + format="PNG", + value=field_a, + field_0=field_b, + ) + assert False, "PNG must reject extra layers" + except ValueError as exc: + assert "single layer" in str(exc).lower() + + +def test_save_multi_channel_gwy_round_trip(): + """A multi-channel GWY save round-trips via the gwy importer.""" + from backend.importers import gwy as gwy_importer + from backend.nodes.save import Save + + rng = np.random.default_rng(53) + primary = DataField( + data=rng.standard_normal((24, 32)).astype(np.float64) * 1e-9, + xreal=4e-6, yreal=3e-6, si_unit_xy="m", si_unit_z="m", + ) + layer2 = DataField( + data=rng.standard_normal((24, 32)).astype(np.float64) * 1e-11, + xreal=4e-6, yreal=3e-6, si_unit_xy="m", si_unit_z="N", + ) + + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "topo" + Save().save( + filename=str(path), + format="GWY", + value=primary, + field_0=layer2, + primary_name="height", + layer_name_0="adhesion", + ) + out_path = path.with_suffix(".gwy") + assert out_path.exists() + + reloaded = gwy_importer.load(out_path) + assert len(reloaded) == 2 + names = gwy_importer.channel_names(out_path) + assert set(names) == {"height", "adhesion"} + # GWY does not guarantee iteration order across channels, so match + # each input by content rather than by position. + assert any(np.allclose(f.data, primary.data) for f in reloaded) + assert any(np.allclose(f.data, layer2.data) for f in reloaded) + for f in reloaded: + assert np.isclose(f.xreal, 4e-6) + assert np.isclose(f.yreal, 3e-6) + + +def test_save_multi_channel_hdf5_round_trip(): + """Multi-channel generic HDF5 round-trips via the hdf5 importer.""" + from backend.importers import hdf5 as hdf5_importer + from backend.nodes.save import Save + + rng = np.random.default_rng(59) + primary = DataField( + data=rng.standard_normal((12, 18)).astype(np.float64) * 1e-7, + xreal=2e-6, yreal=1.5e-6, si_unit_xy="m", si_unit_z="V", + ) + layer2 = DataField( + data=rng.standard_normal((12, 18)).astype(np.float64) * 1e-9, + xreal=2e-6, yreal=1.5e-6, si_unit_xy="m", si_unit_z="A", + ) + + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "stack" + Save().save( + filename=str(path), + format="HDF5", + value=primary, + field_0=layer2, + primary_name="potential", + layer_name_0="current", + ) + out_path = path.with_suffix(".h5") + assert out_path.exists() + + reloaded = hdf5_importer.load(out_path) + assert len(reloaded) == 2 + # Identify the two channels by their unique z-units. + by_unit = {rf.si_unit_z: rf for rf in reloaded} + assert set(by_unit.keys()) == {"V", "A"} + assert np.allclose(by_unit["V"].data, primary.data) + assert np.allclose(by_unit["A"].data, layer2.data) + + +def test_save_multi_channel_hdf5_ergo_round_trip(): + """Multi-channel Ergo-layout HDF5 round-trips via the ergo_hdf5 importer.""" + from backend.importers import ergo_hdf5 as ergo_importer + from backend.nodes.save import Save + + rng = np.random.default_rng(61) + primary = DataField( + data=rng.standard_normal((10, 14)).astype(np.float64) * 1e-9, + xreal=1.5e-6, yreal=1e-6, si_unit_xy="m", si_unit_z="m", + ) + layer2 = DataField( + data=rng.standard_normal((10, 14)).astype(np.float64) * 1e-11, + xreal=1.5e-6, yreal=1e-6, si_unit_xy="m", si_unit_z="N", + ) + + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "topo" + Save().save( + filename=str(path), + format="HDF5 (Ergo)", + value=primary, + field_0=layer2, + primary_name="height", + layer_name_0="adhesion", + ) + out_path = path.with_suffix(".h5") + assert out_path.exists() + + reloaded = ergo_importer.load(out_path) + assert len(reloaded) == 2 + by_unit = {rf.si_unit_z: rf for rf in reloaded} + assert set(by_unit.keys()) == {"m", "N"} + assert np.allclose(by_unit["m"].data, primary.data) + assert np.allclose(by_unit["N"].data, layer2.data) + + +def test_save_gwy_rejects_image_layer(): + """GWY/HDF5 formats must error cleanly on non-DataField layers.""" + from backend.nodes.save import Save + + field = DataField(data=np.zeros((4, 4))) + image = np.zeros((4, 4, 3), dtype=np.uint8) + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "topo" + try: + Save().save( + filename=str(path), + format="GWY", + value=field, + field_0=image, + ) + assert False, "GWY must reject non-DataField layers" + except ValueError as exc: + assert "DataField" in str(exc) or "data field" in str(exc).lower() + + +def test_save_ignores_extra_layers_for_non_stackable_types(): + """Stray field_N kwargs must be ignored when value is a scalar/line/table.""" + from backend.nodes.save import Save + + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "scalar" + # field_0 is connected but should be silently ignored for a FLOAT value. + Save().save( + filename=str(path), + format="TXT", + value=1.25, + field_0=DataField(data=np.zeros((4, 4))), + ) + assert Path(tmpdir, "scalar.txt").read_text(encoding="utf-8").strip() == "1.25" + + def test_tiff_preview_is_still_rgb_uint8(): """The legacy TIFF format for DATA_FIELD must keep producing 8-bit RGB.""" import tifffile diff --git a/tests/node_tests/save_layers.py b/tests/node_tests/save_layers.py deleted file mode 100644 index e5130b7..0000000 --- a/tests/node_tests/save_layers.py +++ /dev/null @@ -1,77 +0,0 @@ -import os -import tempfile - -import numpy as np -import tifffile -from PIL import Image - -from tests.node_tests._shared import make_field - - -def test_save_image(): - from backend.nodes.save_layers import SaveImage - - node = SaveImage() - input_types = SaveImage.INPUT_TYPES() - field_spec = input_types["optional"]["field_0"] - assert field_spec[0] == "DATA_FIELD" - assert field_spec[1]["accepted_types"] == ["IMAGE", "ANNOTATION_SOURCE"] - - field_a = make_field(data=np.random.default_rng(4).random((32, 32))) - field_b = make_field(data=np.random.default_rng(5).random((32, 32))) - annotated = np.zeros((24, 24, 3), dtype=np.uint8) - annotated[..., 0] = 255 - - with tempfile.TemporaryDirectory() as tmpdir: - tiff_path = os.path.join(tmpdir, "out.tiff") - node.save(filename=tiff_path, format="TIFF", field_0=field_a) - assert os.path.exists(tiff_path) - im = Image.open(tiff_path) - assert im.n_frames == 1 - assert np.array(im).shape == (32, 32) - - tiff_path2 = os.path.join(tmpdir, "multi.tiff") - node.save(filename=tiff_path2, format="TIFF", field_0=field_a, field_1=field_b) - im2 = Image.open(tiff_path2) - assert im2.n_frames == 2 - - annotated_tiff = os.path.join(tmpdir, "annotated.tiff") - node.save(filename=annotated_tiff, format="TIFF", field_0=annotated, layer_name_0="annotated overview") - with tifffile.TiffFile(annotated_tiff) as tif: - assert len(tif.pages) == 1 - assert tif.pages[0].description == "annotated overview" - assert tif.pages[0].asarray().shape == annotated.shape - - npz_path = os.path.join(tmpdir, "out.npz") - node.save(filename=npz_path, format="NPZ", field_0=field_a, field_1=annotated, layer_name_0="height map", layer_name_1="annotated-overview") - assert os.path.exists(npz_path) - npz = np.load(npz_path) - assert len(npz.files) == 2 - assert np.allclose(npz["height_map"], field_a.data) - assert np.array_equal(npz["annotated_overview"], annotated) - - wrong_ext = os.path.join(tmpdir, "output.png") - node.save(filename=wrong_ext, format="TIFF", field_0=field_a) - assert os.path.exists(os.path.join(tmpdir, "output.tiff")) - - driven_dir = os.path.join(tmpdir, "nested-output") - node.save(filename="driven_name", directory=driven_dir, format="NPZ", field_0=field_a) - assert os.path.exists(os.path.join(driven_dir, "driven_name.npz")) - - try: - node.save(filename="bad", directory=os.path.join(tmpdir, "looks_like_file.txt"), format="TIFF", field_0=field_a) - assert False, "Should have raised ValueError for file-like directory path" - except ValueError: - pass - - try: - node.save(filename=os.path.join(tmpdir, "empty.tiff"), format="TIFF") - assert False, "Should have raised ValueError" - except ValueError: - pass - - try: - node.save(filename="", format="TIFF", field_0=field_a) - assert False, "Should have raised ValueError" - except ValueError: - pass