From 619885676b66d57c66f667399b04f58e0e9a542e Mon Sep 17 00:00:00 2001 From: Benjamin Gallusser Date: Tue, 29 Jun 2021 16:02:50 +0200 Subject: [PATCH] Expose chunk shape in write nodes --- gunpowder/nodes/hdf5like_write_base.py | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/gunpowder/nodes/hdf5like_write_base.py b/gunpowder/nodes/hdf5like_write_base.py index 11a96f13..9e177689 100644 --- a/gunpowder/nodes/hdf5like_write_base.py +++ b/gunpowder/nodes/hdf5like_write_base.py @@ -41,6 +41,12 @@ class Hdf5LikeWrite(BatchFilter): A dictionary from array keys to datatype (eg. ``np.int8``). If given, arrays are stored using this type. The original arrays within the pipeline remain unchanged. + + chunks (``tuple`` of ``int``, or ``bool``): + + Chunk shape for output datasets. Set to ``True`` for auto-chunking, + set to ``False`` to obtain a chunk equal to the dataset size. + Defaults to ``True``. ''' def __init__( @@ -49,7 +55,8 @@ def __init__( output_dir='.', output_filename='output.hdf', compression_type=None, - dataset_dtypes=None): + dataset_dtypes=None, + chunks=True): self.dataset_names = dataset_names self.output_dir = output_dir @@ -59,6 +66,7 @@ def __init__( self.dataset_dtypes = {} else: self.dataset_dtypes = dataset_dtypes + self.chunks = chunks self.dataset_offsets = {} @@ -149,10 +157,12 @@ def init_datasets(self, batch): offset, voxel_size) dataset = data_file.create_dataset( - name=dataset_name, - shape=data_shape, - compression=self.compression_type, - dtype=dtype) + name=dataset_name, + shape=data_shape, + compression=self.compression_type, + dtype=dtype, + chunks=self.chunks + ) self._set_offset(dataset, offset) self._set_voxel_size(dataset, voxel_size)