diff --git a/aicsimageio/readers/bioformats_reader.py b/aicsimageio/readers/bioformats_reader.py index 1a2bce3e1..04104697c 100644 --- a/aicsimageio/readers/bioformats_reader.py +++ b/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: diff --git a/aicsimageio/tests/readers/test_bioformats_reader.py b/aicsimageio/tests/readers/test_bioformats_reader.py index fa65117dd..804544301 100644 --- a/aicsimageio/tests/readers/test_bioformats_reader.py +++ b/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) diff --git a/scripts/TEST_RESOURCES_HASH.txt b/scripts/TEST_RESOURCES_HASH.txt index f892f4b56..e7a7918e3 100644 --- a/scripts/TEST_RESOURCES_HASH.txt +++ b/scripts/TEST_RESOURCES_HASH.txt @@ -1 +1 @@ -f65e7b1ecd3eccf384e2fb1333fff841bbde06a1d38b9985a613ef06dd7530bb +3a475850cd52c29f4cc047289e770434fe5fe33863742fa03e08623635bf49bd