Map bioformats images with tiles in x-y dimension in dask

aicsimageio/readers/bioformats_reader.py

@@ -75,7 +75,13 @@ class BioformatsReader(Reader):
         see: https://docs.openmicroscopy.org/bio-formats/latest/formats/options.html
         For example: to turn off chunkmap table reading for ND2 files, use
         `options={"nativend2.chunkmap": False}`
-
+    dask_tiles: bool, optional
+        Whether to chunk the bioformats dask array by tiles to easily read sub-regions
+        with numpy-like array indexing
+        Defaults to false and iamges are read by entire planes
+    tile_size: Optional[Tuple[int, int]]
+        Tuple that sets the tile size of y and x axis, respectively
+        By default, it will use optimal values computed by bioformats itself
     Raises
     ------
     exceptions.UnsupportedFileFormatError
@@ -101,6 +107,8 @@ def __init__(
         original_meta: bool = False,
         memoize: Union[int, bool] = 0,
         options: Dict[str, bool] = {},
+        dask_tiles: bool = False,
+        tile_size: Optional[Tuple[int, int]] = None,
     ):
         self._fs, self._path = io_utils.pathlike_to_fs(image, enforce_exists=True)
         # Catch non-local file system
@@ -114,6 +122,8 @@ def __init__(
             "options": options,
             "original_meta": original_meta,
             "memoize": memoize,
+            "dask_tiles": dask_tiles,
+            "tile_size": tile_size,
         }
         try:
             with BioFile(self._path, **self._bf_kwargs) as rdr:  # type: ignore
@@ -247,6 +257,13 @@ class BioFile:
         see: https://docs.openmicroscopy.org/bio-formats/latest/formats/options.html
         For example: to turn off chunkmap table reading for ND2 files, use
         `options={"nativend2.chunkmap": False}`
+    dask_tiles: bool, optional
+        Whether to chunk the bioformats dask array by tiles to easily read sub-regions
+        with numpy-like array indexing
+        Defaults to false and iamges are read by entire planes
+    tile_size: Optional[Tuple[int, int]]
+        Tuple that sets the tile size of y and x axis, respectively
+        By default, it will use optimal values computed by bioformats itself
     """
 
     def __init__(
@@ -258,6 +275,8 @@ def __init__(
         original_meta: bool = False,
         memoize: Union[int, bool] = 0,
         options: Dict[str, bool] = {},
+        dask_tiles: bool = False,
+        tile_size: Optional[Tuple[int, int]] = None,
     ):
         try:
             loci = get_loci()
@@ -299,6 +318,16 @@ def __init__(
         self._lock = Lock()
         self.set_series(series)
 
+        self.dask_tiles = dask_tiles
+        if self.dask_tiles:
+            if tile_size is None:
+                self.tile_size = (
+                    self._r.getOptimalTileHeight(),
+                    self._r.getOptimalTileWidth(),
+                )
+            else:
+                self.tile_size = tile_size
+
     def set_series(self, series: int = 0) -> None:
         self._r.setSeries(series)
         self._core_meta = CoreMeta(
@@ -367,7 +396,12 @@ def to_dask(self, series: Optional[int] = None) -> DaskArrayProxy:
             self._r.setSeries(series)
 
         nt, nc, nz, ny, nx, nrgb = self.core_meta.shape
-        chunks = ((1,) * nt, (1,) * nc, (1,) * nz, ny, nx)
+
+        if self.dask_tiles:
+            chunks = _get_dask_tile_chunks(nt, nc, nz, ny, nx, self.tile_size)
+        else:
+            chunks = ((1,) * nt, (1,) * nc, (1,) * nz, (ny,), (nx,))
+
         if nrgb > 1:
             chunks = chunks + (nrgb,)  # type: ignore
         arr = da.map_blocks(
@@ -481,8 +515,16 @@ def _dask_chunk(self, block_id: Tuple[int, ...]) -> np.ndarray:
         """
         # Our convention is that the final dask array is in the order TCZYX, so
         # block_id will be coming in as (T, C, Z, Y, X).
-        t, c, z, *_ = block_id
-        im = self._get_plane(t, c, z)
+        t, c, z, y, x, *_ = block_id
+
+        if self.dask_tiles:
+            *_, ny, nx, _ = self.core_meta.shape
+            y_slice = _axis_id_to_slice(y, self.tile_size[0], ny)
+            x_slice = _axis_id_to_slice(x, self.tile_size[1], nx)
+            im = self._get_plane(t, c, z, y_slice, x_slice)
+        else:
+            im = self._get_plane(t, c, z)
+
         return im[np.newaxis, np.newaxis, np.newaxis]
 
     _service: Any = None
@@ -517,6 +559,43 @@ def _pixtype2dtype(pixeltype: int, little_endian: bool) -> np.dtype:
     return np.dtype(("<" if little_endian else ">") + fmt2type[pixeltype])
 
 
+def _chunk_by_tile_size(n_px: int, tile_length: int) -> Tuple[int, ...]:
+    n_splits = n_px / tile_length
+    n_full_tiles = np.floor(n_splits)
+
+    if n_splits.is_integer():
+        tile_chunks = (int(tile_length),) * int(n_full_tiles)
+    else:
+        edge_tile = n_px - (n_full_tiles * tile_length)
+        tile_chunks = (int(tile_length),) * int(n_full_tiles) + (int(edge_tile),)
+    return tile_chunks
+
+
+def _get_dask_tile_chunks(
+    nt: int, nc: int, nz: int, ny: int, nx: int, tile_size: Tuple[int, int]
+) -> Tuple[
+    Tuple[int, ...], Tuple[int, ...], Tuple[int, ...], Tuple[int, ...], Tuple[int, ...]
+]:
+    """Returns chunking tuples (length of each chunk in each axis) after tiling.
+    I.e., if nx == 2048 and tile_size == 1024, chunks for x axis will be (1024,1024)"""
+
+    y_tile_size, x_tile_size = tile_size
+
+    y_tiling_chunks = _chunk_by_tile_size(ny, y_tile_size)
+    x_tiling_chunks = _chunk_by_tile_size(nx, x_tile_size)
+
+    return ((1,) * nt, (1,) * nc, (1,) * nz, y_tiling_chunks, x_tiling_chunks)
+
+
+def _axis_id_to_slice(axis_id: int, tile_length: int, n_px: int) -> slice:
+    """Take the axis_id from a dask block_id and create the corresponding
+    tile slice, taking into account edge tiles."""
+    if (axis_id * tile_length) + tile_length <= n_px:
+        return slice(axis_id * tile_length, (axis_id * tile_length) + tile_length)
+    else:
+        return slice(axis_id * tile_length, n_px)
+
+
 def _slice2width(slc: slice, length: int) -> Tuple[int, int]:
     """Convert `slice` object into (start, width)"""
     if slc.stop is not None or slc.start is not None:

aicsimageio/tests/readers/test_bioformats_reader.py

@@ -404,3 +404,54 @@ def test_multi_scene_bioformats_reader(
         second_scene_shape=second_scene_shape,
         second_scene_dtype=np.dtype(np.uint16),
     )
+
+
+@pytest.mark.parametrize(
+    "filename, ",
+    [
+        ("CMU-1-Small-Region.svs"),
+    ],
+)
+def test_bioformats_dask_tiling_shapes(filename: str) -> None:
+    # Construct full filepath
+    uri = get_resource_full_path(filename, LOCAL)
+
+    # Run checks
+    bf_tiled = BioformatsReader(uri, dask_tiles=True)
+    bf_fullplane = BioformatsReader(uri, dask_tiles=False)
+    bf_tiled_set = BioformatsReader(uri, dask_tiles=True, tile_size=(1024, 1024))
+
+    np.testing.assert_array_equal(
+        bf_tiled.dask_data.shape, bf_fullplane.dask_data.shape
+    )
+    np.testing.assert_array_equal(
+        bf_tiled_set.dask_data.shape, bf_fullplane.dask_data.shape
+    )
+    np.testing.assert_array_equal(
+        bf_tiled.dask_data.shape, bf_fullplane.dask_data.shape
+    )
+    np.testing.assert_array_equal(bf_tiled.dask_data.chunksize, (1, 1, 1, 240, 240, 3))
+    np.testing.assert_array_equal(
+        bf_tiled_set.dask_data.chunksize, (1, 1, 1, 1024, 1024, 3)
+    )
+
+
+@pytest.mark.parametrize(
+    "filename, ",
+    [("s_1_t_1_c_10_z_1.ome.tiff"), ("CMU-1-Small-Region.svs")],
+)
+def test_bioformats_dask_tiling_read(filename: str) -> None:
+    # Construct full filepath
+    uri = get_resource_full_path(filename, LOCAL)
+
+    # Run checks
+    bf_tiled = BioformatsReader(uri, dask_tiles=True)
+    bf_tiled_set = BioformatsReader(uri, dask_tiles=True, tile_size=(128, 128))
+    bf_fullplane = BioformatsReader(uri, dask_tiles=False)
+
+    arr_tiled = bf_tiled.dask_data.compute()
+    arr_fullplane = bf_fullplane.dask_data.compute()
+    arr_tiled_set = bf_tiled_set.dask_data.compute()
+
+    np.testing.assert_array_equal(arr_tiled, arr_fullplane)
+    np.testing.assert_array_equal(arr_tiled_set, arr_fullplane)

scripts/TEST_RESOURCES_HASH.txt

@@ -1 +1 @@
-f65e7b1ecd3eccf384e2fb1333fff841bbde06a1d38b9985a613ef06dd7530bb
+3a475850cd52c29f4cc047289e770434fe5fe33863742fa03e08623635bf49bd