Repository: scalableminds/webknossos-libs

Issue:
NIFTI Conversion Performance
We received feedback that the NIFTI conversion is:
1) slow
2) and has memory leaks. 

Researchers at Brown university tried converting their large-scale EM Nifti data and ran into the above issues.


Discussion:

Resolving patch:
diff --git a/wkcuber/convert_nifti.py b/wkcuber/convert_nifti.py
--- a/wkcuber/convert_nifti.py
+++ b/wkcuber/convert_nifti.py
@@ -2,6 +2,7 @@
 import time
 from argparse import ArgumentParser
 from pathlib import Path
+from sklearn.preprocessing import LabelEncoder
 from typing import Tuple, Optional, Union, cast
 
 import nibabel as nib
@@ -40,6 +41,15 @@ def create_parser() -> ArgumentParser:
         default="color",
     )
 
+    parser.add_argument(
+        "--is_segmentation_layer",
+        "-s",
+        help="When converting one layer, signals whether layer is segmentation layer. "
+        "When converting a folder, this option is ignored",
+        default=False,
+        action="store_true",
+    )
+
     parser.add_argument(
         "--dtype", "-d", help="Target datatype (e.g. uint8, uint16).", default="uint8"
     )
@@ -90,13 +100,17 @@ def create_parser() -> ArgumentParser:
 
 
 def to_target_datatype(
-    data: np.ndarray, target_dtype: str, is_probably_binary: bool
+    data: np.ndarray, target_dtype: str, is_segmentation_layer: bool
 ) -> np.ndarray:
-    if is_probably_binary:
-        logging.info(
-            f"Casting directly to {target_dtype}, as input seems to be binary."
+
+    if is_segmentation_layer:
+        original_shape = data.shape
+        label_encoder = LabelEncoder()
+        return (
+            label_encoder.fit_transform(data.ravel())
+            .reshape(original_shape)
+            .astype(np.dtype(target_dtype))
         )
-        return data.astype(np.dtype(target_dtype))
 
     if data.dtype == np.dtype("float32"):
         factor = data.max()
@@ -118,6 +132,7 @@ def convert_nifti(
     layer_name: str,
     dtype: str,
     scale: Tuple[float, ...],
+    is_segmentation_layer: bool = False,
     file_len: int = DEFAULT_WKW_FILE_LEN,
     bbox_to_enforce: BoundingBox = None,
     write_tiff: bool = False,
@@ -137,14 +152,12 @@ def convert_nifti(
 
     cube_data = np.array(source_nifti.get_fdata())
 
-    is_probably_binary = np.unique(cube_data).shape[0] == 2
-    assume_segmentation_layer = False  # Since webKnossos does not support multiple segmention layers, this is hardcoded to False right now.
     max_cell_id_args = (
         {"largest_segment_id": int(np.max(cube_data) + 1)}
-        if assume_segmentation_layer
+        if is_segmentation_layer
         else {}
     )
-    category_type = "segmentation" if assume_segmentation_layer else "color"
+    category_type = "segmentation" if is_segmentation_layer else "color"
     logging.debug(f"Assuming {category_type} as layer type for {layer_name}")
 
     if len(source_nifti.shape) == 3:
@@ -171,8 +184,9 @@ def convert_nifti(
 
     if scale is None:
         scale = tuple(map(float, source_nifti.header["pixdim"][:3]))
+
     logging.info(f"Using scale: {scale}")
-    cube_data = to_target_datatype(cube_data, dtype, is_probably_binary)
+    cube_data = to_target_datatype(cube_data, dtype, is_segmentation_layer)
 
     # everything needs to be padded to
     if bbox_to_enforce is not None:
@@ -273,9 +287,10 @@ def convert_folder_nifti(
                 "color",
                 "uint8",
                 scale,
-                write_tiff,
-                bbox_to_enforce,
-                use_orientation_header,
+                is_segmentation_layer=False,
+                write_tiff=write_tiff,
+                bbox_to_enforce=bbox_to_enforce,
+                use_orientation_header=use_orientation_header,
                 flip_axes=flip_axes,
             )
         elif path == segmentation_path:
@@ -285,9 +300,10 @@ def convert_folder_nifti(
                 "segmentation",
                 "uint8",
                 scale,
-                write_tiff,
-                bbox_to_enforce,
-                use_orientation_header,
+                is_segmentation_layer=True,
+                write_tiff=write_tiff,
+                bbox_to_enforce=bbox_to_enforce,
+                use_orientation_header=use_orientation_header,
                 flip_axes=flip_axes,
             )
         else:
@@ -297,9 +313,10 @@ def convert_folder_nifti(
                 path.stem,
                 "uint8",
                 scale,
-                write_tiff,
-                bbox_to_enforce,
-                use_orientation_header,
+                is_segmentation_layer=False,
+                write_tiff=write_tiff,
+                bbox_to_enforce=bbox_to_enforce,
+                use_orientation_header=use_orientation_header,
                 flip_axes=flip_axes,
             )
 
@@ -340,6 +357,7 @@ def main() -> None:
             Path(args.target_path),
             args.layer_name,
             args.dtype,
+            is_segmentation_layer=args.is_segmentation_layer,
             **conversion_args,
         )
 


Test patch:
diff --git a/tests/test_nifti_conversion.py b/tests/test_nifti_conversion.py
new file mode 100644
--- /dev/null
+++ b/tests/test_nifti_conversion.py
@@ -0,0 +1,34 @@
+import numpy as np
+from wkcuber.convert_nifti import to_target_datatype
+
+
+def test_to_target_datatype_segmentation_layer():
+    data = np.array(
+        [
+            [
+                [[0, 0, 1, 5, 27], [3200, 3, 2, 1, 5], [0, 3, 2, 1, 5]],
+                [[3200, 3, 2, 1, 5], [0, 0, 1, 5, 27], [6, 6, 6, 6, 6]],
+            ]
+        ]
+    ).astype("uint32")
+
+    labelled_data = np.array(
+        [
+            [
+                [[0, 0, 1, 4, 6], [7, 3, 2, 1, 4], [0, 3, 2, 1, 4]],
+                [[7, 3, 2, 1, 4], [0, 0, 1, 4, 6], [5, 5, 5, 5, 5]],
+            ]
+        ]
+    )
+
+    target_dtype = np.dtype("uint32")
+
+    converted_segmentation = to_target_datatype(
+        data, target_dtype, is_segmentation_layer=True
+    )
+    converted_no_segmentation = to_target_datatype(
+        data, target_dtype, is_segmentation_layer=False
+    )
+
+    assert np.all(converted_segmentation == labelled_data)
+    assert np.all(converted_no_segmentation == data)