Merge pull request #53 from thewtex/compress-encode-multiscale

Author: thewtex

server/scripts/compress_encode.py

@@ -4,16 +4,19 @@
 import itk
 from numcodecs import Blosc
 import zarr
+from pathlib import Path
+import numpy as np
 
 def compress_encode(input_filepath,
                     output_directory,
+                    multiscale=True,
                     chunk_size=64,
                     cname='zstd',
                     clevel=5,
                     shuffle=True):
     image = itk.imread(input_filepath)
     image_da = itk.xarray_from_image(image)
-    dataset_name = input_filepath
+    dataset_name = str(Path(input_filepath))
     image_ds = image_da.to_dataset(name=dataset_name)
 
     store_name = output_directory
@@ -31,6 +34,37 @@ def compress_encode(input_filepath,
 
     zarr.consolidate_metadata(store)
 
+    if multiscale:
+        # multi-resolution pyramid
+        pyramid = [image_da]
+        reduced = image
+        while not np.all(np.array(itk.size(reduced)) < 64):
+            scale = len(pyramid)
+            shrink_factors = [2**scale]*3
+            reduced = itk.bin_shrink_image_filter(image, shrink_factors=shrink_factors)
+            reduced_da = itk.xarray_from_image(reduced)
+            pyramid.append(reduced_da)
+
+        pyramid_group_paths = [""]
+        for scale in range(1, len(pyramid)):
+            pyramid_group_paths.append('scale_{0}'.format(scale))
+
+        for scale in range(1, len(pyramid)):
+            ds = pyramid[scale].to_dataset(name=dataset_name)
+            ds.to_zarr(store,
+                       mode='w',
+                       group=pyramid_group_paths[scale],
+                       compute=True,
+                       encoding={dataset_name: {'chunks': [chunk_size]*3, 'compressor': compressor}})
+
+        # Re-consolidate entire dataset
+        zarr.consolidate_metadata(store)
+        for scale in range(1, len(pyramid)):
+            store = zarr.DirectoryStore(str(Path(store_name) / pyramid_group_paths[scale]))
+            # Also consolidate the metadata on the pyramid scales so they can be used independently
+            zarr.consolidate_metadata(store)
+
+
 if __name__ == '__main__':
     parser = argparse.ArgumentParser('Convert and encode a medical image file in a compressed Zarr directory store.')
     parser.add_argument('input_filepath', help='Path to input image file, e.g. a NIFTI file.')
@@ -40,11 +74,13 @@ def compress_encode(input_filepath,
     parser.add_argument('--chunk-size', default=64, type=int, help='Compression chunk size along one dimension.')
     parser.add_argument('--cname', default='zstd', help='Base compression codec.')
     parser.add_argument('--clevel', default=5, type=int, help='Compression level.')
+    parser.add_argument('--no-multi-scale', action='store_true', help='Do not generate a multi-scale pyramid.')
 
     args = parser.parse_args()
 
     compress_encode(args.input_filepath,
                     args.output_directory,
+                    multiscale=not args.no_multi_scale,
                     chunk_size=args.chunk_size,
                     cname=args.cname,
                     clevel=args.clevel,

server/scripts/requirements.txt

@@ -0,0 +1,3 @@
+itk
+xarray
+zarr