diff --git a/brainglobe_utils/image_io/save.py b/brainglobe_utils/image_io/save.py index 163c943..1a35069 100644 --- a/brainglobe_utils/image_io/save.py +++ b/brainglobe_utils/image_io/save.py @@ -1,4 +1,5 @@ import warnings +from pathlib import Path import nrrd import numpy as np @@ -104,3 +105,38 @@ def to_nrrd(img_volume, dest_path): """ dest_path = str(dest_path) nrrd.write(dest_path, img_volume) + + +def save_any(img_volume, dest_path): + """ + Save the image volume (numpy array) to the given file path, using the save + function matching its file extension. + + Parameters + ---------- + img_volume : np.ndarray + The image to be saved. + + dest_path : str or pathlib.Path + The file path to save the image to. + Supports directories (will save a sequence of tiffs), .tif, .tiff, + .nrrd and .nii. + """ + dest_path = Path(dest_path) + + if dest_path.is_dir(): + to_tiffs(img_volume, str(dest_path / "image")) + + elif dest_path.suffix == ".tif" or dest_path.suffix == ".tiff": + to_tiff(img_volume, str(dest_path)) + + elif dest_path.suffix == ".nrrd": + to_nrrd(img_volume, str(dest_path)) + + elif dest_path.suffix == ".nii": + to_nii(img_volume, str(dest_path)) + + else: + raise NotImplementedError( + f"Could not guess data type for path {dest_path}" + ) diff --git a/tests/tests/test_image_io.py b/tests/tests/test_image_io.py index 1ff37a5..d460dfd 100644 --- a/tests/tests/test_image_io.py +++ b/tests/tests/test_image_io.py @@ -29,46 +29,18 @@ def image_array(request, array_2d, array_3d): @pytest.fixture() -def shuffled_txt_path(tmp_path, array_3d): +def txt_path(tmp_path, array_3d): """ Return the path to a text file containing the paths of a series of 2D tiffs - in a random order + in order """ txt_path = tmp_path / "imgs_file.txt" - write_tiff_sequence_with_txt_file(txt_path, array_3d) - - # Shuffle paths in the text file into a random order - with open(txt_path, "r+") as f: - tiff_paths = f.read().splitlines() - random.Random(4).shuffle(tiff_paths) - f.seek(0) - f.writelines(line + "\n" for line in tiff_paths) - f.truncate() - - return txt_path - - -def write_tiff_sequence_with_txt_file(txt_path, image_array): - """ - Write an image array to a series of tiffs, and write a text file - containing all the tiff file paths in order (one per line). - - The tiff sequence will be saved to a sub-folder inside the same folder - as the text file. - - Parameters - ---------- - txt_path : pathlib.Path - Filepath of text file to create - image_array : np.ndarray - Image to write as sequence of tiffs - """ directory = txt_path.parent # Write tiff sequence to sub-folder sub_dir = directory / "sub" sub_dir.mkdir() - save.to_tiffs(image_array, str(sub_dir / "image_array")) + save.to_tiffs(array_3d, str(sub_dir / "image")) # Write txt file containing all tiff file paths (one per line) tiff_paths = sorted(sub_dir.iterdir()) @@ -76,33 +48,24 @@ def write_tiff_sequence_with_txt_file(txt_path, image_array): "\n".join([str(sub_dir / fname) for fname in tiff_paths]) ) + return txt_path -def save_any(file_path, image_array): - """ - Save image_array to given file path, using the save function matching - its file extension - Parameters - ---------- - file_path : pathlib.Path - File path of image to save - image_array : np.ndarray - Image to save +@pytest.fixture() +def shuffled_txt_path(txt_path): """ - if file_path.is_dir(): - save.to_tiffs(image_array, str(file_path / "image_array")) - - elif file_path.suffix == ".txt": - write_tiff_sequence_with_txt_file(file_path, image_array) - - elif file_path.suffix == ".tif" or file_path.suffix == ".tiff": - save.to_tiff(image_array, str(file_path)) - - elif file_path.suffix == ".nrrd": - save.to_nrrd(image_array, str(file_path)) + Return the path to a text file containing the paths of a series of 2D tiffs + in a random order + """ + # Shuffle paths in the text file into a random order + with open(txt_path, "r+") as f: + tiff_paths = f.read().splitlines() + random.Random(4).shuffle(tiff_paths) + f.seek(0) + f.writelines(line + "\n" for line in tiff_paths) + f.truncate() - elif file_path.suffix == ".nii": - save.to_nii(image_array, str(file_path), scale=(1, 1, 1)) + return txt_path def test_tiff_io(tmp_path, image_array): @@ -110,9 +73,9 @@ def test_tiff_io(tmp_path, image_array): Test that a 2D/3D tiff can be written and read correctly """ dest_path = tmp_path / "image_array.tiff" - save_any(dest_path, image_array) + save.save_any(image_array, dest_path) + reloaded = load.load_any(str(dest_path)) - reloaded = load.load_img_stack(str(dest_path), 1, 1, 1) assert (reloaded == image_array).all() @@ -127,11 +90,14 @@ def test_3d_tiff_scaling( Test that a 3D tiff is scaled correctly on loading """ dest_path = tmp_path / "image_array.tiff" - save_any(dest_path, array_3d) - - reloaded = load.load_img_stack( - str(dest_path), x_scaling_factor, y_scaling_factor, z_scaling_factor + save.save_any(array_3d, dest_path) + reloaded = load.load_any( + str(dest_path), + x_scaling_factor=x_scaling_factor, + y_scaling_factor=y_scaling_factor, + z_scaling_factor=z_scaling_factor, ) + assert reloaded.shape[0] == array_3d.shape[0] * z_scaling_factor assert reloaded.shape[1] == array_3d.shape[1] * y_scaling_factor assert reloaded.shape[2] == array_3d.shape[2] * x_scaling_factor @@ -149,10 +115,10 @@ def test_tiff_sequence_io(tmp_path, array_3d, load_parallel): Test that a 3D image can be written and read correctly as a sequence of 2D tiffs (with or without parallel loading) """ - save_any(tmp_path, array_3d) - reloaded_array = load.load_from_folder( - str(tmp_path), 1, 1, 1, load_parallel=load_parallel - ) + save.save_any(array_3d, tmp_path) + assert len(list(tmp_path.glob("*.tif"))) == array_3d.shape[0] + + reloaded_array = load.load_any(str(tmp_path), load_parallel=load_parallel) assert (reloaded_array == array_3d).all() @@ -166,9 +132,12 @@ def test_tiff_sequence_scaling( """ Test that a tiff sequence is scaled correctly on loading """ - save_any(tmp_path, array_3d) - reloaded_array = load.load_from_folder( - str(tmp_path), x_scaling_factor, y_scaling_factor, z_scaling_factor + save.save_any(array_3d, tmp_path) + reloaded_array = load.load_any( + str(tmp_path), + x_scaling_factor=x_scaling_factor, + y_scaling_factor=y_scaling_factor, + z_scaling_factor=z_scaling_factor, ) assert reloaded_array.shape[0] == array_3d.shape[0] * z_scaling_factor @@ -176,19 +145,6 @@ def test_tiff_sequence_scaling( assert reloaded_array.shape[2] == array_3d.shape[2] * x_scaling_factor -def test_load_img_sequence_from_txt(tmp_path, array_3d): - """ - Test that a tiff sequence can be loaded from a text file containing an - ordered list of the tiff file paths (one per line) - """ - img_sequence_file = tmp_path / "imgs_file.txt" - save_any(img_sequence_file, array_3d) - - # Load image from paths in text file - reloaded_array = load.load_img_sequence(str(img_sequence_file), 1, 1, 1) - assert (reloaded_array == array_3d).all() - - @pytest.mark.parametrize( "sort", [True, False], @@ -209,11 +165,16 @@ def test_sort_img_sequence_from_txt(shuffled_txt_path, array_3d, sort): def test_nii_io(tmp_path, array_3d): """ - Test that a 3D image can be written and read correctly as nii + Test that a 3D image can be written and read correctly as nii with scale + (keeping it as a nifty object with no numpy conversion on loading) """ - nii_path = tmp_path / "test_array.nii" - save_any(nii_path, array_3d) - assert (load.load_nii(str(nii_path)).get_fdata() == array_3d).all() + nii_path = str(tmp_path / "test_array.nii") + scale = (5, 5, 5) + save.to_nii(array_3d, nii_path, scale=scale) + reloaded = load.load_nii(nii_path) + + assert (reloaded.get_fdata() == array_3d).all() + assert reloaded.header.get_zooms() == scale def test_nii_read_to_numpy(tmp_path, array_3d): @@ -221,42 +182,41 @@ def test_nii_read_to_numpy(tmp_path, array_3d): Test that conversion of loaded nii image to an in-memory numpy array works """ nii_path = tmp_path / "test_array.nii" - save_any(nii_path, array_3d) + save.save_any(array_3d, nii_path) + reloaded_array = load.load_any(str(nii_path), as_numpy=True) - reloaded_array = load.load_nii(str(nii_path), as_array=True, as_numpy=True) assert (reloaded_array == array_3d).all() -def test_nrrd_io(tmp_path, array_3d): - """ - Test that a 3D image can be written and read correctly as nrrd - """ - nrrd_path = tmp_path / "test_array.nrrd" - save_any(nrrd_path, array_3d) - assert (load.load_nrrd(str(nrrd_path)) == array_3d).all() - - @pytest.mark.parametrize( "file_name", [ "test_array.tiff", "test_array.tif", - "test_array.txt", "test_array.nrrd", "test_array.nii", pytest.param("", id="dir of tiffs"), ], ) -def test_load_any(tmp_path, array_3d, file_name): +def test_save_and_load_any(tmp_path, array_3d, file_name): """ - Test that load_any can read all required image file types + Test that save_any/load_any can write/read all required image + file types. """ src_path = tmp_path / file_name - save_any(src_path, array_3d) + save.save_any(array_3d, src_path) assert (load.load_any(str(src_path)) == array_3d).all() +def test_load_any_txt(txt_path, array_3d): + """ + Test that load_any can read a tiff sequence from a text file containing an + ordered list of the tiff file paths (one per line) + """ + assert (load.load_any(str(txt_path)) == array_3d).all() + + def test_load_any_error(tmp_path): """ Test that load_any throws an error for an unknown file extension @@ -265,6 +225,14 @@ def test_load_any_error(tmp_path): load.load_any(str(tmp_path / "test.unknown")) +def test_save_any_error(tmp_path, array_3d): + """ + Test that save_any throws an error for an unknown file extension + """ + with pytest.raises(NotImplementedError): + save.save_any(array_3d, str(tmp_path / "test.unknown")) + + def test_scale_z(array_3d): """ Test that a 3D image can be scaled in z by float and integer values @@ -273,22 +241,24 @@ def test_scale_z(array_3d): assert utils.scale_z(array_3d, 2).shape[0] == array_3d.shape[0] * 2 -@pytest.mark.parametrize( - "file_name", - [ - "test_array.txt", - pytest.param("", id="dir of tiffs"), - ], -) -def test_image_size(tmp_path, array_3d, file_name): +def test_image_size_dir(tmp_path, array_3d): """ - Test that image size can be detected from a directory of 2D tiffs, or - a text file containing the paths of a sequence of 2D tiffs + Test that image size can be detected from a directory of 2D tiffs """ - file_path = tmp_path / file_name - save_any(file_path, array_3d) + save.save_any(array_3d, tmp_path) - image_shape = load.get_size_image_from_file_paths(str(file_path)) + image_shape = load.get_size_image_from_file_paths(str(tmp_path)) + assert image_shape["x"] == array_3d.shape[2] + assert image_shape["y"] == array_3d.shape[1] + assert image_shape["z"] == array_3d.shape[0] + + +def test_image_size_txt(txt_path, array_3d): + """ + Test that image size can be detected from a text file containing the paths + of a sequence of 2D tiffs + """ + image_shape = load.get_size_image_from_file_paths(str(txt_path)) assert image_shape["x"] == array_3d.shape[2] assert image_shape["y"] == array_3d.shape[1] assert image_shape["z"] == array_3d.shape[0]