diff --git a/looptrace/conversion_to_zarr.py b/looptrace/conversion_to_zarr.py index 40c4881e..370fe4af 100644 --- a/looptrace/conversion_to_zarr.py +++ b/looptrace/conversion_to_zarr.py @@ -6,6 +6,7 @@ EMBL Heidelberg """ +from collections import OrderedDict import os from pathlib import * from typing import * @@ -17,6 +18,7 @@ from looptrace import image_io, nd2io, ArrayDimensionalityError from looptrace.ImageHandler import ImageHandler from looptrace.integer_naming import get_fov_name_short +from looptrace.nd2io import CHANNEL_COUNT_KEY from looptrace.voxel_stack import VoxelSize @@ -24,18 +26,28 @@ def workflow(n_pos: int, input_folders: Iterable[Path], output_folder: Path) -> for fov_index in tqdm.tqdm(range(int(n_pos))): imgs = [] for f in input_folders: - folder_imgs, _, folder_metadata = nd2io.stack_nd2_to_dask(f, fov_index=fov_index) + folder_imgs, _, sample_file_metadata = nd2io.stack_nd2_to_dask(f, fov_index=fov_index) imgs.append(folder_imgs[0]) imgs = da.concatenate(imgs, axis=0) - if len(imgs.shape) == 5: - chunks = (1, 1, 1, imgs.shape[-2], imgs.shape[-1]), # 1 chunk per xy-plane (z-slice) + + num_channels: int = sample_file_metadata[CHANNEL_COUNT_KEY] + exp_num_dim: int + if num_channels == 1: + exp_num_dim = 4 + imgs = imgs.reshape((imgs.shape[0], 1) + imgs.shape[1:]) # Create channel axis. + elif num_channels > 1: + exp_num_dim = 5 else: + raise RuntimeError(f"Channel count isn't positive: {num_channels}") + if len(imgs.shape) != exp_num_dim: raise ArrayDimensionalityError( - f"Expected a 5D array to write to ZARR, but got {len(imgs.shape)}D; shape: {imgs.shape}" + f"Expected a {exp_num_dim}-D array to write to ZARR, but got {len(imgs.shape)}-D; shape: {imgs.shape}" ) - voxel_size: VoxelSize = folder_metadata["voxel_size"] + + voxel_size: VoxelSize = sample_file_metadata["voxel_size"] + chunks = (1, 1, 1, imgs.shape[-2], imgs.shape[-1]), # 1 chunk per xy-plane (z-slice) - # TODO: why is it justified to use just the last folder_metadata value (associated with a + # TODO: why is it justified to use just the last sample_file_metadata value (associated with a # single f in input_folders) in a function call where the concatenation of values from # all input_folders is being passed to .zarr creation? # See: https://github.com/gerlichlab/looptrace/issues/118 @@ -46,7 +58,7 @@ def workflow(n_pos: int, input_folders: Iterable[Path], output_folder: Path) -> shape = imgs.shape, dtype = np.uint16, chunks = chunks, - metadata = folder_metadata, + metadata = sample_file_metadata, voxel_size = voxel_size, ) n_t = imgs.shape[0] diff --git a/looptrace/nd2io.py b/looptrace/nd2io.py index 5f4b6f1a..188ffb56 100644 --- a/looptrace/nd2io.py +++ b/looptrace/nd2io.py @@ -27,7 +27,8 @@ ] -_AXIS_SIZES_KEY = "axis_sizes" +AXIS_SIZES_KEY = "axis_sizes" +CHANNEL_COUNT_KEY = "channelCount" class EmptyImagesError(Exception): @@ -77,7 +78,8 @@ def parse_nd2_metadata(image_file: str) -> Mapping[str, Any]: metadata = {} with nd2.ND2File(image_file) as sample: metadata["voxel_size"] = parse_voxel_size(sample) - metadata[_AXIS_SIZES_KEY] = sample.sizes + metadata[AXIS_SIZES_KEY] = sample.sizes + metadata[CHANNEL_COUNT_KEY] = getattr(sample.attributes, CHANNEL_COUNT_KEY) microscope = sample.metadata.channels[0].microscope metadata['microscope'] = { 'objectiveMagnification': microscope.objectiveMagnification, @@ -169,7 +171,7 @@ def _shift_axes_of_stacked_array_from_nd2( arr: da.Array, metadata: Mapping[str, Any], ) -> Result[da.Array, str]: - match list(metadata[_AXIS_SIZES_KEY].keys()): + match list(metadata[AXIS_SIZES_KEY].keys()): case ["Z", "C", "Y", "X"]: return Result.Ok(da.moveaxis(arr, -4, -3)) case ["Z", "Y", "X"]: diff --git a/tests/test__stack_nd2_to_dask.py b/tests/test__stack_nd2_to_dask.py index b90dfd84..55eba368 100644 --- a/tests/test__stack_nd2_to_dask.py +++ b/tests/test__stack_nd2_to_dask.py @@ -12,7 +12,7 @@ __email__ = "vincent.reuter@imba.oeaw.ac.at" from looptrace.nd2io import * -from looptrace.nd2io import _AXIS_SIZES_KEY +from looptrace.nd2io import AXIS_SIZES_KEY from looptrace.integer_naming import get_fov_names_N POSITION_PREFIX = "Point000" @@ -179,7 +179,7 @@ def test_underscore_prefixed_and_or_non_nd2_files_are_skipped_and_good_ones_have unique_fields_of_view.add(p) # Patch the metadata parser to be a no-op, the ND2 reader to be context manager-like, # and dask call to be identity. - with mock.patch("looptrace.nd2io.parse_nd2_metadata", return_value={_AXIS_SIZES_KEY: OrderedDict((dim, 0) for dim in ["Z", "C", "Y", "X"])}), \ + with mock.patch("looptrace.nd2io.parse_nd2_metadata", return_value={AXIS_SIZES_KEY: OrderedDict((dim, 0) for dim in ["Z", "C", "Y", "X"])}), \ mock.patch("looptrace.nd2io.nd2.ND2File", side_effect=lambda *_, **__: mocked_nd2_handle()) as mock_nd2_read, \ mock.patch("looptrace.nd2io.da.stack", side_effect=lambda arrs: arrs), \ mock.patch("looptrace.nd2io.da.moveaxis", side_effect=lambda _1, _2, _3: mock.Mock(shape=None)):