128 lines
4.4 KiB
Python
128 lines
4.4 KiB
Python
from __future__ import annotations
|
|
import numpy as np
|
|
from backend.node_registry import register_node
|
|
from backend.execution_context import emit_warning
|
|
from backend.data_types import (
|
|
COLORMAPS,
|
|
DataField,
|
|
ImageData,
|
|
_apply_annotation_overlay_from_context,
|
|
_annotation_context_from_image,
|
|
image_to_uint8,
|
|
normalize_font_spec,
|
|
resolve_colormap_input,
|
|
)
|
|
|
|
|
|
@register_node(display_name="Annotations")
|
|
class Annotations:
|
|
_broadcast_warning_fn = None
|
|
_current_node_id: str = ""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"input": ("ANNOTATION_SOURCE", {
|
|
"label": "Input",
|
|
"accepted_types": ["DATA_FIELD", "IMAGE"],
|
|
}),
|
|
"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 = ("ANNOTATION_SOURCE",)
|
|
RETURN_NAMES = ("Output",)
|
|
FUNCTION = "render"
|
|
|
|
DESCRIPTION = (
|
|
"Attach optional publication-style annotations to a DATA_FIELD without flattening the raw data, "
|
|
"or annotate an IMAGE that carries viewport metadata from View3D."
|
|
)
|
|
|
|
def render(
|
|
self,
|
|
input,
|
|
colormap: str,
|
|
show_scale_bar: bool,
|
|
show_color_map: bool,
|
|
text_size: float = 1.0,
|
|
colormap_map=None,
|
|
font=None,
|
|
) -> tuple:
|
|
annotation_spec = {
|
|
"kind": "annotation",
|
|
"show_scale_bar": bool(show_scale_bar),
|
|
"show_color_map": bool(show_color_map),
|
|
"text_size": float(np.clip(text_size, 6.0, 96.0)) if np.isfinite(text_size) else 14.0,
|
|
"font": normalize_font_spec(font),
|
|
}
|
|
|
|
if isinstance(input, DataField):
|
|
resolved_colormap = resolve_colormap_input(
|
|
colormap,
|
|
colormap_input=colormap_map,
|
|
inherited=input.colormap,
|
|
default="gray",
|
|
)
|
|
out = input.replace(
|
|
colormap=resolved_colormap,
|
|
overlays=[
|
|
*input.overlays,
|
|
annotation_spec,
|
|
],
|
|
)
|
|
return (out,)
|
|
|
|
context = _annotation_context_from_image(input)
|
|
if context is None:
|
|
self._send_warning(
|
|
"Annotations image input has no scale metadata, so scale bar and color-map legend cannot be added."
|
|
)
|
|
return (ImageData(image_to_uint8(input)),)
|
|
|
|
resolved_colormap = resolve_colormap_input(
|
|
colormap,
|
|
colormap_input=colormap_map,
|
|
inherited=context.get("colormap"),
|
|
default="gray",
|
|
)
|
|
context["colormap"] = resolved_colormap
|
|
missing_features = []
|
|
xreal = context.get("xreal")
|
|
if bool(show_scale_bar) and not (isinstance(xreal, (int, float)) and np.isfinite(float(xreal)) and float(xreal) > 0 and str(context.get("si_unit_xy", "")).strip()):
|
|
missing_features.append("scale bar")
|
|
if bool(show_color_map):
|
|
legend_values = (context.get("legend_min"), context.get("legend_mid"), context.get("legend_max"))
|
|
has_legend_values = all(
|
|
isinstance(value, (int, float)) and np.isfinite(float(value))
|
|
for value in legend_values
|
|
)
|
|
if not (has_legend_values and str(context.get("legend_unit", "")).strip()):
|
|
missing_features.append("color-map legend")
|
|
if missing_features:
|
|
self._send_warning(
|
|
f"Annotations image input is missing metadata for: {', '.join(missing_features)}."
|
|
)
|
|
annotated = _apply_annotation_overlay_from_context(
|
|
image_to_uint8(input),
|
|
context,
|
|
annotation_spec,
|
|
)
|
|
return (ImageData(annotated, metadata={"annotation_context": context}),)
|
|
|
|
def _send_warning(self, message: str):
|
|
emit_warning(message)
|