From 91fee9f518a32747a199e7c2a1364cbe5a5eb190 Mon Sep 17 00:00:00 2001
From: Tom Herold <tom.herold@scalableminds.com>
Date: Fri, 7 Feb 2020 11:50:44 +0100
Subject: [PATCH] Enable Compression for NIFTI Conversion Script (#182)

* Update convert_nifti.py

* fix formatting

* fixed compression
---
 .vscode/settings.json    |  3 +++
 wkcuber/convert_nifti.py | 27 +++++++++++++++++++++++++--
 2 files changed, 28 insertions(+), 2 deletions(-)
 create mode 100644 .vscode/settings.json

diff --git a/.vscode/settings.json b/.vscode/settings.json
new file mode 100644
index 000000000..684dee7ba
--- /dev/null
+++ b/.vscode/settings.json
@@ -0,0 +1,3 @@
+{
+    "python.pythonPath": "/usr/local/Caskroom/miniconda/base/bin/python"
+}
\ No newline at end of file
diff --git a/wkcuber/convert_nifti.py b/wkcuber/convert_nifti.py
index a35213b38..a03eae2d3 100644
--- a/wkcuber/convert_nifti.py
+++ b/wkcuber/convert_nifti.py
@@ -71,9 +71,18 @@ def to_target_datatype(data: np.ndarray, target_dtype) -> np.ndarray:
     return (data / factor).astype(np.dtype(target_dtype))
 
 
-def convert_nifti(source_nifti_path, target_path, layer_name, dtype, scale, mag=1):
+def convert_nifti(
+    source_nifti_path, target_path, layer_name, dtype, scale, mag=1, file_len=256
+):
     target_wkw_info = WkwDatasetInfo(
-        str(target_path.resolve()), layer_name, mag, wkw.Header(np.dtype(dtype))
+        str(target_path.resolve()),
+        layer_name,
+        mag,
+        wkw.Header(
+            np.dtype(dtype),
+            block_type=wkw.Header.BLOCK_TYPE_LZ4HC,
+            file_len=file_len // 32,
+        ),
     )
     ensure_wkw(target_wkw_info)
 
@@ -106,6 +115,20 @@ def convert_nifti(source_nifti_path, target_path, layer_name, dtype, scale, mag=
             scale = tuple(map(float, source_nifti.header["pixdim"][:3]))
 
         cube_data = to_target_datatype(cube_data, dtype)
+
+        # Writing wkw compressed requires files of shape (file_len, file_len, file_len)
+        # Pad data accordingly
+        padding_offset = file_len - np.array(cube_data.shape[1:4]) % file_len
+        cube_data = np.pad(
+            cube_data,
+            (
+                (0, 0),
+                (0, int(padding_offset[0])),
+                (0, int(padding_offset[1])),
+                (0, int(padding_offset[2])),
+            ),
+        )
+
         target_wkw.write(offset, cube_data)
 
     logging.debug(