78 lines
3.1 KiB
Python
78 lines
3.1 KiB
Python
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
|