Merge pull request #56 from thewtex/compress-encode-path

OpenImaging · Feb 11, 2021 · dc015bb · dc015bb
2 parents a38b626 + 3c33020
commit dc015bb
Showing 1 changed file with 13 additions and 8 deletions.
diff --git a/server/scripts/compress_encode.py b/server/scripts/compress_encode.py
@@ -28,7 +28,7 @@ def zip_zchunkstore(zip_file, url=None):
 
     JSON-serializable reference description.
     """
-    rfs = {}
+    zchunkstore = {}
     with zipfile.ZipFile(zip_file) as zf:
         if zf.compression != 0:
             raise RuntimeError("Compressed zip's are not supported.")
@@ -39,7 +39,6 @@ def zip_zchunkstore(zip_file, url=None):
         if url is not None:
             data_url = url
 
-        zchunkstore = {}
         for info in zf.infolist():
             name_bytes = len(info.filename.encode("utf-8"))
             offset = info.header_offset + 30 + name_bytes
@@ -50,7 +49,7 @@ def zip_zchunkstore(zip_file, url=None):
             else:
                 zchunkstore[info.filename] = [data_url, offset, size]
 
-        return zchunkstore
+    return zchunkstore
 
 
 def compress_encode(input_filepath,
@@ -60,10 +59,11 @@ def compress_encode(input_filepath,
                     cname='zstd',
                     clevel=5,
                     shuffle=True,
-                    zip_chunk_store=True):
+                    zip_chunk_store=True,
+                    zip_store_url=None):
     image = itk.imread(input_filepath)
     image_da = itk.xarray_from_image(image)
-    dataset_name = str(Path(input_filepath))
+    dataset_name = Path(input_filepath).stem
     image_ds = image_da.to_dataset(name=dataset_name)
 
     store_name = output_directory
@@ -88,7 +88,7 @@ def compress_encode(input_filepath,
         reduced = image
         while not np.all(np.array(itk.size(reduced)) < 64):
             scale = len(pyramid)
-            shrink_factors = [2]*3
+            shrink_factors = [2]*image.GetImageDimension()
             for i, s in enumerate(itk.size(reduced)):
                 if s < 4:
                     shrink_factors[i] = 1
@@ -120,7 +120,10 @@ def compress_encode(input_filepath,
         zip_store_path = str(Path(output_directory)) + '.zip'
         with zarr.storage.ZipStore(zip_store_path, mode='w', compression=0) as zip_store:
             zarr.copy_store(store, zip_store)
-        zchunkstore = zip_zchunkstore(zip_store_path)
+        store_url = None
+        if zip_store_url:
+            store_url = zip_store_url
+        zchunkstore = zip_zchunkstore(zip_store_path, store_url)
         with open(zip_store_path + '.zchunkstore', 'w') as fp:
             json.dump(zchunkstore, fp)
 
@@ -136,6 +139,7 @@ def compress_encode(input_filepath,
     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.')
     parser.add_argument('--no-zip-chunk-store', action='store_true', help='Do not generate a zip file and corresponding chunk store.')
+    parser.add_argument('--zip-store-url', help='URL where the generated zip file will be hosted.')
 
     args = parser.parse_args()
 
@@ -146,4 +150,5 @@ def compress_encode(input_filepath,
                     cname=args.cname,
                     clevel=args.clevel,
                     shuffle=not args.no_shuffle,
-                    zip_chunk_store=not args.no_zip_chunk_store)
+                    zip_chunk_store=not args.no_zip_chunk_store,
+                    zip_store_url=args.zip_store_url)