diff --git a/qim3d/io/saving.py b/qim3d/io/saving.py index 64a6a1ca5adfd9efab8b7f8d2cb7ebb4f5b25b0f..b163c3126705962e0c437a50a86aeb28ed043254 100644 --- a/qim3d/io/saving.py +++ b/qim3d/io/saving.py @@ -24,12 +24,14 @@ Example: import datetime import os +import dask.array as da import h5py import nibabel as nib import numpy as np import PIL import pydicom import tifffile +import zarr from pydicom.dataset import FileDataset, FileMetaDataset from pydicom.uid import UID @@ -234,6 +236,21 @@ class DataSaver: ds.PixelData = data_bytes ds.save_as(path) + + def save_to_zarr(self, path, data): + """ Saves a Dask array to a Zarr array on disk. + + Args: + path (str): The path to the Zarr array on disk. + data (dask.array): The Dask array to be saved to disk. + + Returns: + zarr.core.Array: The Zarr array saved on disk. + """ + assert isinstance(data, da.Array), 'data must be a dask array' + + # 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)) def save_PIL(self, path, data): @@ -330,6 +347,8 @@ class DataSaver: return self.save_vol(path, data) elif path.endswith((".dcm",".DCM")): return self.save_dicom(path, data) + elif path.endswith((".zarr")): + return self.save_to_zarr(path, data) elif path.endswith((".jpeg",".jpg", ".png")): return self.save_PIL(path, data) else: