diff --git a/brainglobe_utils/image_io/load.py b/brainglobe_utils/image_io/load.py index fb87b1a..33e5acf 100644 --- a/brainglobe_utils/image_io/load.py +++ b/brainglobe_utils/image_io/load.py @@ -15,6 +15,7 @@ get_num_processes, get_sorted_file_paths, ) +from brainglobe_utils.image_io.utils import ImageIOLoadException from .utils import check_mem, scale_z @@ -85,6 +86,15 @@ def load_any( ------- np.ndarray The loaded brain. + + Raises + ------ + ImageIOLoadException + If there was an issue loading the image with image_io. + + See Also + ------ + image_io.utils.ImageIOLoadException """ src_path = Path(src_path) @@ -169,7 +179,8 @@ def load_img_stack( Parameters ---------- stack_path : str or pathlib.Path - The path of the image to be loaded. + The path of the image to be loaded. Note that only 3D tiffs are + supported. x_scaling_factor : float The scaling of the brain along the x dimension (applied on loading @@ -192,11 +203,19 @@ def load_img_stack( ------- np.ndarray The loaded brain array. + + Raises + ------ + ImageIOLoadException + If attempt to load a 2D tiff. """ stack_path = Path(stack_path) logging.debug(f"Loading: {stack_path}") stack = tifffile.imread(stack_path) + if stack.ndim != 3: + raise ImageIOLoadException(error_type="2D tiff") + # Downsampled plane by plane because the 3D downsampling in scipy etc # uses too much RAM @@ -217,11 +236,9 @@ def load_img_stack( logging.debug("Converting downsampled stack to array") stack = np.array(downsampled_stack) - if stack.ndim == 3: - # stack = np.rollaxis(stack, 0, 3) - if z_scaling_factor != 1: - logging.debug("Downsampling stack in Z") - stack = scale_z(stack, z_scaling_factor) + if z_scaling_factor != 1: + logging.debug("Downsampling stack in Z") + stack = scale_z(stack, z_scaling_factor) return stack @@ -278,7 +295,8 @@ def load_from_folder( Parameters ---------- src_folder : str or pathlib.Path - The source folder containing tiff files. + The source folder containing tiff files. Note that this folder must + contain at least 2 tiffs, and all tiff images must have the same shape. x_scaling_factor : float, optional The scaling of the brain along the x dimension (applied on loading @@ -310,6 +328,12 @@ def load_from_folder( ------- np.ndarray The loaded and scaled brain. + + Raises + ------ + ImageIOLoadException + If attempt to load a directory containing only a single tiff, or a + sequence of tiffs that have different shapes. """ paths = get_sorted_file_paths(src_folder, file_extension=file_extension) @@ -342,7 +366,8 @@ def load_img_sequence( ---------- img_sequence_file_path : str or pathlib.Path The path to the file containing the ordered list of image paths (one - per line). + per line). Note that this file must contain at least 2 paths, and all + referenced images must have the same shape. x_scaling_factor : float, optional The scaling of the brain along the x dimension (applied on loading @@ -374,6 +399,12 @@ def load_img_sequence( ------- np.ndarray The loaded and scaled brain. + + Raises + ------ + ImageIOLoadException + If attempt to load a txt file containing only a single path, or a + sequence of paths that load images with different shapes. """ img_sequence_file_path = Path(img_sequence_file_path) with open(img_sequence_file_path, "r") as in_file: @@ -408,7 +439,8 @@ def load_image_series( Parameters ---------- paths : list of str or list of pathlib.Path - Ordered list of image paths. + Ordered list of image paths. Must contain at least 2 paths, and all + referenced images must have the same shape. x_scaling_factor : float, optional The scaling of the brain along the x dimension (applied on loading @@ -436,7 +468,18 @@ def load_image_series( ------- np.ndarray The loaded and scaled brain. + + Raises + ------ + ImageIOLoadException + If attempt to load a single path, or a sequence of paths that load + images with different shapes. """ + # Throw an error if there's only one image to load - should be an image + # series, so at least 2 paths. + if len(paths) == 1: + raise ImageIOLoadException("single_tiff") + if load_parallel: img = threaded_load_from_sequence( paths, @@ -473,7 +516,8 @@ def threaded_load_from_sequence( Parameters ---------- paths_sequence : list of str or list of pathlib.Path - The sorted list of the planes paths on the filesystem. + The sorted list of the planes paths on the filesystem. All planes + must have the same shape. x_scaling_factor : float, optional The scaling of the brain along the x dimension (applied on loading @@ -495,6 +539,11 @@ def threaded_load_from_sequence( ------- np.ndarray The loaded and scaled brain. + + Raises + ------ + ImageIOLoadException + If attempt to load a sequence of images with different shapes. """ stacks = [] n_processes = get_num_processes(min_free_cpu_cores=n_free_cpus) @@ -524,7 +573,17 @@ def threaded_load_from_sequence( anti_aliasing=anti_aliasing, ) stacks.append(process) - stack = np.dstack([s.result() for s in stacks]) + + stack_shapes = set() + for i in range(len(stacks)): + stacks[i] = stacks[i].result() + stack_shapes.add(stacks[i].shape[0:2]) + + # Raise an error if the x/y shape of all stacks aren't the same + if len(stack_shapes) > 1: + raise ImageIOLoadException("sequence_shape") + + stack = np.dstack(stacks) return stack @@ -542,7 +601,8 @@ def load_from_paths_sequence( Parameters ---------- paths_sequence : list of str or list of pathlib.Path - The sorted list of the planes paths on the filesystem. + The sorted list of the planes paths on the filesystem. All planes + must have the same shape. x_scaling_factor : float, optional The scaling of the brain along the x dimension (applied on loading @@ -561,6 +621,11 @@ def load_from_paths_sequence( ------- np.ndarray The loaded and scaled brain. + + Raises + ------ + ImageIOLoadException + If attempt to load a sequence of images with different shapes. """ for i, p in enumerate( tqdm(paths_sequence, desc="Loading images", unit="plane") @@ -590,14 +655,18 @@ def load_from_paths_sequence( preserve_range=True, anti_aliasing=anti_aliasing, ) + + # Raise an error if the shapes of the images aren't the same + if not volume[:, :, i].shape == img.shape: + raise ImageIOLoadException("sequence_shape") volume[:, :, i] = img return volume def get_size_image_from_file_paths(file_path, file_extension="tif"): """ - Returns the size of an image (which is a list of 2D files), without loading - the whole image. + Returns the size of an image (which is a list of 2D tiff files), + without loading the whole image. Parameters ---------- @@ -621,7 +690,7 @@ def get_size_image_from_file_paths(file_path, file_extension="tif"): logging.debug( "Loading file: {} to check raw image size" "".format(img_paths[0]) ) - image_0 = load_any(img_paths[0]) + image_0 = tifffile.imread(img_paths[0]) y_shape, x_shape = image_0.shape image_shape = {"x": x_shape, "y": y_shape, "z": z_shape} diff --git a/brainglobe_utils/image_io/utils.py b/brainglobe_utils/image_io/utils.py index 4f1b970..07e5ed5 100644 --- a/brainglobe_utils/image_io/utils.py +++ b/brainglobe_utils/image_io/utils.py @@ -3,7 +3,53 @@ class ImageIOLoadException(Exception): - pass + """ + Custom exception class for errors found loading images with + image_io.load. + + Alerts the user of: loading a directory containing only a single .tiff, + loading a single 2D .tiff, loading an image sequence where all 2D images + don't have the same shape, lack of memory to complete loading. + + Set the error message to self.message to read during testing. + """ + + def __init__(self, error_type=None, total_size=None, free_mem=None): + if error_type == "single_tiff": + self.message = ( + "Attempted to load directory containing " + "a single .tiff file. If the .tiff file " + "is 3D please pass the full path with " + "filename. Single 2D .tiff file input is " + "not supported." + ) + + elif error_type == "2D tiff": + self.message = "Single 2D .tiff file input is not supported." + + elif error_type == "sequence_shape": + self.message = ( + "Attempted to load an image sequence where individual 2D " + "images did not have the same shape. Please ensure all image " + "files contain the same number of pixels." + ) + + elif error_type == "memory": + self.message = ( + "Not enough memory on the system to complete " + "loading operation." + ) + if total_size is not None and free_mem is not None: + self.message += ( + f" Needed {total_size}, only {free_mem} " f"available." + ) + + else: + self.message = ( + "File failed to load with brainglobe_utils.image_io." + ) + + super().__init__(self.message) def check_mem(img_byte_size, n_imgs): @@ -24,15 +70,14 @@ def check_mem(img_byte_size, n_imgs): Raises ------ - BrainLoadException + ImageIOLoadException If not enough memory is available. """ total_size = img_byte_size * n_imgs free_mem = psutil.virtual_memory().available if total_size >= free_mem: raise ImageIOLoadException( - "Not enough memory on the system to complete loading operation" - "Needed {}, only {} available.".format(total_size, free_mem) + error_type="memory", total_size=total_size, free_mem=free_mem ) diff --git a/tests/tests/test_image_io.py b/tests/tests/test_image_io.py index 992a74a..e720cd5 100644 --- a/tests/tests/test_image_io.py +++ b/tests/tests/test_image_io.py @@ -1,6 +1,8 @@ import random +from collections import namedtuple import numpy as np +import psutil import pytest from brainglobe_utils.image_io import load, save, utils @@ -19,15 +21,6 @@ def array_3d(array_2d): return volume -@pytest.fixture(params=["2D", "3D"]) -def image_array(request, array_2d, array_3d): - """Create both a 2D and 3D array of 32-bit integers""" - if request.param == "2D": - return array_2d - else: - return array_3d - - @pytest.fixture() def txt_path(tmp_path, array_3d): """ @@ -69,9 +62,9 @@ def shuffled_txt_path(txt_path): @pytest.mark.parametrize("use_path", [True, False], ids=["Path", "String"]) -def test_tiff_io(tmp_path, image_array, use_path): +def test_tiff_io(tmp_path, array_3d, use_path): """ - Test that a 2D/3D tiff can be written and read correctly, using string + Test that a 3D tiff can be written and read correctly, using string or pathlib.Path input. """ filename = "image_array.tiff" @@ -80,10 +73,10 @@ def test_tiff_io(tmp_path, image_array, use_path): else: dest_path = str(tmp_path / filename) - save.to_tiff(image_array, dest_path) + save.to_tiff(array_3d, dest_path) reloaded = load.load_img_stack(dest_path, 1, 1, 1) - assert (reloaded == image_array).all() + assert (reloaded == array_3d).all() @pytest.mark.parametrize( @@ -140,6 +133,17 @@ def test_tiff_sequence_io(tmp_path, array_3d, load_parallel, use_path): assert (reloaded_array == array_3d).all() +def test_2d_tiff(tmp_path, array_2d): + """ + Test that an error is thrown when loading a single 2D tiff + """ + image_path = tmp_path / "image.tif" + save.to_tiff(array_2d, image_path) + + with pytest.raises(utils.ImageIOLoadException): + load.load_any(image_path) + + @pytest.mark.parametrize( "x_scaling_factor, y_scaling_factor, z_scaling_factor", [(1, 1, 1), (0.5, 0.5, 1), (0.25, 0.25, 0.25)], @@ -163,6 +167,36 @@ def test_tiff_sequence_scaling( assert reloaded_array.shape[2] == array_3d.shape[2] * x_scaling_factor +def test_tiff_sequence_one_tiff(tmp_path): + """ + Test that an error is thrown when loading a directory containing a + single tiff via load_any + """ + save.to_tiff(np.ones((3, 3)), tmp_path / "image.tif") + + with pytest.raises(utils.ImageIOLoadException): + load.load_any(tmp_path) + + +@pytest.mark.parametrize( + "load_parallel", + [ + pytest.param(True, id="parallel loading"), + pytest.param(False, id="no parallel loading"), + ], +) +def test_tiff_sequence_diff_shape(tmp_path, array_3d, load_parallel): + """ + Test that an error is thrown when trying to load a tiff sequence where + individual 2D tiffs have different shapes + """ + save.to_tiff(np.ones((2, 2)), tmp_path / "image_1.tif") + save.to_tiff(np.ones((3, 3)), tmp_path / "image_2.tif") + + with pytest.raises(utils.ImageIOLoadException): + load.load_any(tmp_path, load_parallel=load_parallel) + + @pytest.mark.parametrize("use_path", [True, False], ids=["Path", "String"]) def test_load_img_sequence_from_txt(txt_path, array_3d, use_path): """ @@ -325,3 +359,20 @@ def test_image_size_txt(txt_path, array_3d): 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_memory_error(monkeypatch): + """ + Test that check_mem throws an error when there's not enough memory + available. + """ + + # Use monkeypatch to always return a set value for the available memory. + def mock_memory(): + VirtualMemory = namedtuple("VirtualMemory", "available") + return VirtualMemory(500) + + monkeypatch.setattr(psutil, "virtual_memory", mock_memory) + + with pytest.raises(utils.ImageIOLoadException): + utils.check_mem(8, 1000)