remove dtype argument and infer dtype from input data

Author: valentin-pinkau

tests/test_downsampling.py

@@ -8,7 +8,7 @@
 )
 import wkw
 from wkcuber.utils import WkwDatasetInfo, open_wkw
-from wkcuber.downsampling import _mode
+from wkcuber.downsampling import _mode, non_linear_filter_3d
 import shutil
 
 WKW_CUBE_SIZE = 1024
@@ -34,18 +34,43 @@ def test_downsample_cube():
     assert buffer[0, 0, 1] == 1
     assert np.all(output[:, :, :] == np.arange(0, CUBE_EDGE_LEN, 2))
 
+
 def test_downsample_mode():
 
-    a = np.array([[1, 3, 4, 2, 2, 7],
-                  [5, 2, 2, 1, 4, 1],
-                  [3, 3, 2, 2, 1, 1]])
+    a = np.array([[1, 3, 4, 2, 2, 7], [5, 2, 2, 1, 4, 1], [3, 3, 2, 2, 1, 1]])
 
     result = _mode(a)
     expected_result = np.array([1, 3, 2, 2, 1, 1])
 
     assert np.all(result == expected_result)
 
 
+def test_downsample_median():
+
+    a = np.array([[1, 3, 4, 2, 2, 7], [5, 2, 2, 1, 4, 1], [3, 3, 2, 2, 1, 1]])
+
+    result = np.median(a, axis=0)
+    expected_result = np.array([3, 3, 2, 2, 2, 1])
+
+    assert np.all(result == expected_result)
+
+
+def test_non_linear_filter_reshape():
+    a = np.array([[[1, 3], [1, 4]], [[4, 2], [3, 1]]], dtype=np.uint8)
+
+    a_filtered = non_linear_filter_3d(a, [2, 2, 2], _mode)
+    assert a_filtered.dtype == np.uint8
+    expected_result = [1]
+    assert np.all(expected_result == a_filtered)
+
+    a = np.array([[[1, 3], [1, 4]], [[4, 3], [3, 1]]], np.uint32)
+
+    a_filtered = non_linear_filter_3d(a, [2, 2, 1], _mode)
+    assert a_filtered.dtype == np.uint32
+    expected_result = [1, 3]
+    assert np.all(expected_result == a_filtered)
+
+
 def test_cube_addresses():
     addresses = cube_addresses(source_info)
     assert len(addresses) == 5 * 5 * 1

wkcuber/downsampling.py

@@ -60,20 +60,13 @@ def create_parser():
         default="default",
     )
 
-    parser.add_argument(
-        "--dtype",
-        "-d",
-        help="Target datatype (e.g. uint8, uint16, uint32)",
-        default="uint8",
-    )
-
     parser.add_argument(
         "--from_mag",
         "--from",
         "-f",
         help="Resolution to base downsampling on",
         type=str,
-        default='1',
+        default="1",
     )
 
     # Either provide the maximum resolution to be downsampled OR a specific, anisotropic magnification.
@@ -202,7 +195,7 @@ def downsample_cube_job(
                     source_wkw.header.file_len * source_wkw.header.block_len
                 )
                 shape = (num_channels,) + (wkw_cubelength,) * 3
-                file_buffer = np.zeros(shape, target_wkw_info.dtype)
+                file_buffer = np.zeros(shape, target_wkw.header.voxel_type)
                 tile_length = cube_edge_len
                 tile_count_per_dim = wkw_cubelength // tile_length
                 assert (
@@ -343,10 +336,18 @@ def _mode(x):
     shape = list(sort.shape)
     shape[axis] = 1
     # Create a boolean array along strides of unique values
-    strides = np.concatenate([np.zeros(shape=shape, dtype='bool'),
-                                 np.diff(sort, axis=axis) == 0,
-                                 np.zeros(shape=shape, dtype='bool')],
-                                axis=axis).transpose(transpose).ravel()
+    strides = (
+        np.concatenate(
+            [
+                np.zeros(shape=shape, dtype="bool"),
+                np.diff(sort, axis=axis) == 0,
+                np.zeros(shape=shape, dtype="bool"),
+            ],
+            axis=axis,
+        )
+        .transpose(transpose)
+        .ravel()
+    )
     # Count the stride lengths
     counts = np.cumsum(strides)
     counts[~strides] = np.concatenate([[0], np.diff(counts[~strides])])
@@ -392,7 +393,6 @@ def downsample_mag(
     layer_name,
     source_mag: Mag,
     target_mag: Mag,
-    dtype="uint8",
     interpolation_mode="default",
     cube_edge_len=DEFAULT_EDGE_LEN,
     jobs=1,
@@ -407,12 +407,8 @@ def downsample_mag(
     else:
         interpolation_mode = InterpolationModes[interpolation_mode.upper()]
 
-    source_wkw_info = WkwDatasetInfo(
-        path, layer_name, dtype, source_mag.to_layer_name()
-    )
-    target_wkw_info = WkwDatasetInfo(
-        path, layer_name, dtype, target_mag.to_layer_name()
-    )
+    source_wkw_info = WkwDatasetInfo(path, layer_name, None, source_mag.to_layer_name())
+    target_wkw_info = WkwDatasetInfo(path, layer_name, None, target_mag.to_layer_name())
     downsample(
         source_wkw_info,
         target_wkw_info,
@@ -430,7 +426,6 @@ def downsample_mags(
     layer_name,
     from_mag: Mag,
     max_mag: Mag,
-    dtype,
     interpolation_mode,
     cube_edge_len,
     jobs,
@@ -444,7 +439,6 @@ def downsample_mags(
             layer_name,
             source_mag,
             target_mag,
-            dtype,
             interpolation_mode,
             cube_edge_len,
             jobs,
@@ -469,7 +463,6 @@ def downsample_mags(
             args.layer_name,
             from_mag,
             anisotropic_target_mag,
-            args.dtype,
             args.interpolation_mode,
             args.buffer_cube_size,
             args.jobs,
@@ -481,7 +474,6 @@ def downsample_mags(
             args.layer_name,
             from_mag,
             max_mag,
-            args.dtype,
             args.interpolation_mode,
             args.buffer_cube_size,
             args.jobs,

wkcuber/utils.py

@@ -21,7 +21,10 @@
 
 
 def _open_wkw(info, **kwargs):
-    header = wkw.Header(np.dtype(info.dtype), **kwargs)
+    if info.dtype is not None:
+        header = wkw.Header(np.dtype(info.dtype), **kwargs)
+    else:
+        header = None
     ds = wkw.Dataset.open(
         path.join(info.dataset_path, info.layer_name, str(info.mag)), header
     )