130 lines
3.9 KiB
Python
130 lines
3.9 KiB
Python
"""
|
|
Exporter for IMAGE values (numpy arrays, ImageData annotation sources).
|
|
|
|
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
|
|
|
|
from backend.data_types import image_to_uint8
|
|
from backend.exporters._base import FormatSpec
|
|
|
|
accepted_types: tuple[str, ...] = ("IMAGE", "ANNOTATION_SOURCE")
|
|
|
|
FORMATS: dict[str, FormatSpec] = {
|
|
"PNG": FormatSpec(ext=".png", round_trip=False, label="PNG"),
|
|
"TIFF": FormatSpec(ext=".tiff", round_trip=False, label="TIFF"),
|
|
"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."
|
|
)
|
|
|
|
if format_name == "PNG":
|
|
from PIL import Image
|
|
Image.fromarray(image_to_uint8(np.asarray(value))).save(str(path))
|
|
return
|
|
if format_name == "TIFF":
|
|
_save_tiff(path, layers, names)
|
|
return
|
|
if format_name == "NPZ":
|
|
_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
|