Refactor ImageIOLoadException

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
 
@@ -197,6 +198,9 @@ def load_img_stack(
     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 +221,10 @@ 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)
+    # stack = np.rollaxis(stack, 0, 3)
+    if z_scaling_factor != 1:
+        logging.debug("Downsampling stack in Z")
+        stack = scale_z(stack, z_scaling_factor)
     return stack
 
 
@@ -437,6 +440,11 @@ def load_image_series(
     np.ndarray
         The loaded and scaled brain.
     """
+    # 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,
@@ -524,7 +532,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
 
 
@@ -590,14 +608,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 +643,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}

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):
@@ -31,8 +77,7 @@
     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
         )
 
 

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)