diff --git a/qim3d/io/saving.py b/qim3d/io/saving.py index dda53eb54558850571de7eb7c6231deb3597aa1f..9ae736258afc09bcb6f4d50f8bd278af684f21f4 100644 --- a/qim3d/io/saving.py +++ b/qim3d/io/saving.py @@ -24,7 +24,6 @@ Example: import datetime import os - import dask.array as da import h5py import nibabel as nib @@ -269,16 +268,23 @@ class DataSaver: zarr.core.Array: The Zarr array saved on disk. """ - # Old version that is using dask - - # assert isinstance(data, da.Array), 'data must be a dask array' + if isinstance(data, da.Array): + # If the data is a Dask array, save using dask + if self.chunk_shape: + log.info("Rechunking data to shape %s", self.chunk_shape) + data = data.rechunk(self.chunk_shape) + log.info("Saving Dask array to Zarr array on disk") + da.to_zarr(data, path, overwrite=self.replace) - # # forces compute when saving to zarr - # da.to_zarr(data, path, compute=True, overwrite=self.replace, compressor=zarr.Blosc(cname='zstd', clevel=3, shuffle=2)) - zarr_array = zarr.open( - path, mode="w", shape=data.shape, chunks=self.chunk_shape, dtype=data.dtype - ) - zarr_array[:] = data + else: + zarr_array = zarr.open( + path, + mode="w", + shape=data.shape, + chunks=self.chunk_shape, + dtype=data.dtype, + ) + zarr_array[:] = data def save_PIL(self, path, data): """Save data to a PIL file to the given path.