82 lines
4.5 KiB
Python
82 lines
4.5 KiB
Python
import numpy as np
|
|
from backend.data_types import DataField, ImageData, datafield_to_uint8, render_datafield_preview
|
|
from backend.execution_context import active_node, execution_callbacks
|
|
|
|
|
|
def test_annotations():
|
|
from backend.nodes.annotations import Annotations
|
|
from backend.nodes.font import Font
|
|
|
|
node = Annotations()
|
|
font_node = Font()
|
|
annotation_input = Annotations.INPUT_TYPES()["required"]["input"]
|
|
assert annotation_input[0] == "ANNOTATION_SOURCE"
|
|
assert annotation_input[1]["accepted_types"] == ["DATA_FIELD", "IMAGE"]
|
|
|
|
warnings = []
|
|
field = DataField(
|
|
data=np.linspace(0.0, 1.0, 64 * 64, dtype=np.float64).reshape(64, 64),
|
|
xreal=1e-6, yreal=1e-6, si_unit_xy="m", si_unit_z="V", colormap="viridis",
|
|
)
|
|
|
|
base = datafield_to_uint8(field, "viridis")
|
|
plain_preview = render_datafield_preview(field, "viridis")
|
|
assert np.array_equal(plain_preview, base)
|
|
|
|
with execution_callbacks(warning=lambda nid, msg: warnings.append(msg)), active_node("test"):
|
|
plain_field, = node.render(input=field, colormap="auto", show_scale_bar=False, show_color_map=False)
|
|
assert isinstance(plain_field, DataField)
|
|
assert np.array_equal(plain_field.data, field.data)
|
|
assert plain_field.colormap == "viridis"
|
|
assert plain_field.overlays[-1]["kind"] == "annotation"
|
|
plain = render_datafield_preview(plain_field, plain_field.colormap)
|
|
assert plain.shape == base.shape
|
|
assert np.array_equal(plain, base)
|
|
|
|
with_scale_field, = node.render(input=field, colormap="auto", show_scale_bar=True, show_color_map=False)
|
|
with_scale = render_datafield_preview(with_scale_field, with_scale_field.colormap)
|
|
assert with_scale.shape == base.shape
|
|
assert not np.array_equal(with_scale, base)
|
|
|
|
with_legend_field, = node.render(input=field, colormap="auto", show_scale_bar=False, show_color_map=True)
|
|
with_legend = render_datafield_preview(with_legend_field, with_legend_field.colormap)
|
|
assert with_legend.shape[0] == base.shape[0]
|
|
assert with_legend.shape[1] > base.shape[1]
|
|
assert with_legend.shape[2] == 3
|
|
|
|
larger_legend_field, = node.render(input=field, colormap="auto", show_scale_bar=False, show_color_map=True, text_size=28.0)
|
|
larger_legend_text = render_datafield_preview(larger_legend_field, larger_legend_field.colormap)
|
|
assert larger_legend_text.shape[0] == with_legend.shape[0]
|
|
assert larger_legend_text.shape[1] > with_legend.shape[1]
|
|
assert not np.array_equal(larger_legend_text, with_legend)
|
|
|
|
annotation_font, = font_node.build("Arial")
|
|
with_font_field, = node.render(input=field, colormap="auto", show_scale_bar=False, show_color_map=True, text_size=28.0, font=annotation_font)
|
|
assert with_font_field.overlays[-1]["font"] == {"family": "Arial", "path": ""}
|
|
with_font = render_datafield_preview(with_font_field, with_font_field.colormap)
|
|
assert with_font.shape[1] > with_legend.shape[1]
|
|
|
|
with_both_field, = node.render(input=field, colormap="auto", show_scale_bar=True, show_color_map=True)
|
|
with_both = render_datafield_preview(with_both_field, with_both_field.colormap)
|
|
assert with_both.shape == with_legend.shape
|
|
assert not np.array_equal(with_both[:, :base.shape[1]], base)
|
|
|
|
viewport_image = ImageData(
|
|
np.zeros((48, 64, 3), dtype=np.uint8),
|
|
metadata={"annotation_context": {"xreal": 2e-6, "si_unit_xy": "m", "legend_min": -1.5, "legend_mid": 0.0, "legend_max": 1.5, "legend_unit": "V", "colormap": "viridis"}},
|
|
)
|
|
annotated_image, = node.render(input=viewport_image, colormap="auto", show_scale_bar=True, show_color_map=True, text_size=18.0)
|
|
assert isinstance(annotated_image, ImageData)
|
|
assert annotated_image.shape[0] == viewport_image.shape[0]
|
|
assert annotated_image.shape[1] > viewport_image.shape[1]
|
|
assert annotated_image.metadata["annotation_context"]["legend_unit"] == "V"
|
|
assert warnings == []
|
|
|
|
plain_image = ImageData(np.zeros((32, 40, 3), dtype=np.uint8))
|
|
passthrough_image, = node.render(input=plain_image, colormap="auto", show_scale_bar=True, show_color_map=True, text_size=18.0)
|
|
assert isinstance(passthrough_image, ImageData)
|
|
assert passthrough_image.shape == plain_image.shape
|
|
assert np.array_equal(np.asarray(passthrough_image), np.asarray(plain_image))
|
|
assert len(warnings) == 1
|
|
assert "no scale metadata" in warnings[0]
|