Enable memory mapping of 3d tiff files where possible

brainglobe_utils/IO/image/load.py

@@ -88,7 +88,7 @@ def load_any(
 
     Returns
     -------
-    np.ndarray
+    np.ndarray or Dask array
         The loaded brain.
 
     Raises
@@ -181,7 +181,7 @@ def load_img_stack(
 
     Returns
     -------
-    np.ndarray
+    np.ndarray or Dask array
         The loaded brain array.
 
     Raises
@@ -191,7 +191,14 @@ def load_img_stack(
     """
     stack_path = Path(stack_path)
     logging.debug(f"Loading: {stack_path}")
-    stack = tifffile.imread(stack_path)
+    try:
+        stack = tifffile.memmap(stack_path, mode="r")
+    except ValueError:
+        try:
+            store = tifffile.imread(stack_path, aszarr=True)
+            stack = da.from_zarr(store)
+        except (ModuleNotFoundError, TypeError):
+            stack = tifffile.imread(stack_path)
 
     if stack.ndim != 3:
         raise ImageIOLoadException(error_type="2D tiff")
@@ -733,8 +740,9 @@ def read_z_stack(path):
     """
     Reads z-stack, lazily, if possible.
 
-    If it's a text file or folder with 2D tiff files use dask to read lazily,
-    otherwise it's a single file tiff stack and is read into memory.
+    If it's a text file or folder with 2D tiff files use dask to read lazily.
+    Otherwise, it's a single file tiff stack and is memory-mapped if possible,
+    otherwise read into memory.
 
     :param path: Filename of text file listing 2D tiffs, folder of 2D tiffs,
         or single file tiff z-stack.
@@ -760,7 +768,14 @@ def read_z_stack(path):
                     "Assume z,y,x"
                 )
 
-        return tifffile.imread(path)
+        try:
+            return tifffile.memmap(path, mode="r")
+        except ValueError:
+            try:
+                store = tifffile.imread(path, aszarr=True)
+                return da.from_zarr(store)
+            except (ModuleNotFoundError, TypeError):
+                return tifffile.imread(path)
 
     return read_with_dask(path)
 

tests/tests/test_IO/test_image_io.py

@@ -98,6 +98,58 @@ def test_tiff_io(tmp_path, array_3d, use_path):
     np.testing.assert_array_equal(reloaded, array_3d)
 
 
+@pytest.mark.parametrize("memmap", [True, False])
+def test_tiff_scale_io(tmp_path, memmap):
+    """
+    Test that a 3D tiff can be read and scaled correctly.
+    """
+    dest_path = tmp_path / "image_array.tiff"
+
+    # compression disables ability to memmap tiff file
+    compression = None if memmap else "zlib"
+    tifffile.imwrite(
+        str(dest_path),
+        np.ones((10, 10, 10)) * 5,
+        metadata={"axes": "ZYX"},
+        compression=compression,
+    )
+
+    reloaded = load.load_img_stack(dest_path, 0.5, 1, 1)
+    np.testing.assert_array_equal(reloaded, np.ones((10, 10, 5)) * 5)
+
+    reloaded = load.load_img_stack(dest_path, 1, 0.5, 1)
+    np.testing.assert_array_equal(reloaded, np.ones((10, 5, 10)) * 5)
+
+    reloaded = load.load_img_stack(dest_path, 1, 1, 0.5)
+    np.testing.assert_array_equal(reloaded, np.ones((5, 10, 10)) * 5)
+
+
+@pytest.mark.parametrize("memmap", [True, False])
+def test_3d_tiff_load_img_stack_memmap_io(tmp_path, array_3d, memmap):
+    """
+    Test that a 3D tiff file can be properly memmapped or loaded using
+    load_img_stack.
+    """
+    filename = str(tmp_path / "image_array.tiff")
+    # compression disables ability to memmap tiff file
+    compression = None if memmap else "zlib"
+
+    tifffile.imwrite(
+        filename, array_3d, metadata={"axes": "ZYX"}, compression=compression
+    )
+
+    reloaded = load.load_img_stack(filename, 1, 1, 1)
+    np.testing.assert_array_equal(reloaded, array_3d)
+
+    if memmap:
+        assert isinstance(reloaded, np.memmap)
+    else:
+        from dask import array as da
+
+        assert not isinstance(reloaded, np.memmap)
+        assert isinstance(reloaded, (np.ndarray, da.Array))
+
+
 @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)],
@@ -419,6 +471,32 @@ def test_read_z_stack_with_missing_metadata(
         mock_debug.assert_called_once()
 
 
+@pytest.mark.parametrize("memmap", [True, False])
+def test_3d_tiff_read_z_stack_memmap_io(tmp_path, array_3d, memmap):
+    """
+    Test that a 3D tiff file can be properly memmapped or loaded using
+    read_z_stack.
+    """
+    filename = str(tmp_path / "image_array.tiff")
+    # compression disables ability to memmap tiff file
+    compression = None if memmap else "zlib"
+
+    tifffile.imwrite(
+        filename, array_3d, metadata={"axes": "ZYX"}, compression=compression
+    )
+
+    reloaded = load.read_z_stack(filename)
+    np.testing.assert_array_equal(reloaded, array_3d)
+
+    if memmap:
+        assert isinstance(reloaded, np.memmap)
+    else:
+        from dask import array as da
+
+        assert not isinstance(reloaded, np.memmap)
+        assert isinstance(reloaded, (np.ndarray, da.Array))
+
+
 def test_get_size_image_with_missing_metadata(
     array3d_as_tiff_stack_with_missing_metadata,
 ):