Implement anisotropic downsampling (#50)

* implement anisotropic downsampling (fixes #49)

* adapt tests to refactored downsampling

* disable dm3 test due to data being unavailable

* change format of anisotropic_target_mag arg

Author: philippotto

.circleci/config.yml

@@ -64,27 +64,27 @@ jobs:
             [ -d testoutput/tiff/color/1 ]
             [ $(find testoutput/tiff/color/1 -mindepth 3 -name "*.wkw" | wc -l) -eq 1 ]
 
-      - run:
-          name: Test DM3 cubing
-          command: |
-            set -x
-            mkdir -p testdata/dm3
-            mkdir -p testoutput/dm3
-            wget http://www.loci.wisc.edu/files/software/data/dnasample1.zip
-            unzip -d testdata/dm3 dnasample1.zip dnasample1.dm3
-            docker run \
-              -v "${PWD}/testdata:/testdata" \
-              -v "${PWD}/testoutput:/testoutput" \
-              --rm \
-              scalableminds/webknossos-cuber:${CIRCLE_BUILD_NUM} \
-              wkcuber.cubing \
-              --verbose \
-              --jobs 1 \
-              --layer_name color \
-              /testdata/dm3 /testoutput/dm3
-            [ -d testoutput/dm3/color ]
-            [ -d testoutput/dm3/color/1 ]
-            [ $(find testoutput/dm3/color/1 -mindepth 3 -name "*.wkw" | wc -l) -eq 16 ]
+      # - run:
+      #     name: Test DM3 cubing
+      #     command: |
+      #       set -x
+      #       mkdir -p testdata/dm3
+      #       mkdir -p testoutput/dm3
+      #       wget http://www.loci.wisc.edu/files/software/data/dnasample1.zip
+      #       unzip -d testdata/dm3 dnasample1.zip dnasample1.dm3
+      #       docker run \
+      #         -v "${PWD}/testdata:/testdata" \
+      #         -v "${PWD}/testoutput:/testoutput" \
+      #         --rm \
+      #         scalableminds/webknossos-cuber:${CIRCLE_BUILD_NUM} \
+      #         wkcuber.cubing \
+      #         --verbose \
+      #         --jobs 1 \
+      #         --layer_name color \
+      #         /testdata/dm3 /testoutput/dm3
+      #       [ -d testoutput/dm3/color ]
+      #       [ -d testoutput/dm3/color/1 ]
+      #       [ $(find testoutput/dm3/color/1 -mindepth 3 -name "*.wkw" | wc -l) -eq 16 ]
 
       - run:
           name: Test tile cubing

tests/test_downsampling.py

@@ -26,7 +26,7 @@ def test_downsample_cube():
     buffer = np.zeros((CUBE_EDGE_LEN,) * 3, dtype=np.uint8)
     buffer[:, :, :] = np.arange(0, CUBE_EDGE_LEN)
 
-    output = downsample_cube(buffer, 2, InterpolationModes.MEDIAN)
+    output = downsample_cube(buffer, (2, 2, 2), InterpolationModes.MEDIAN)
 
     assert output.shape == (CUBE_EDGE_LEN // 2,) * 3
     assert buffer[0, 0, 0] == 0
@@ -59,7 +59,7 @@ def downsample_test_helper(use_compress):
     downsample_cube_job(
         source_info,
         target_info,
-        2,
+        (2, 2, 2),
         InterpolationModes.MAX,
         CUBE_EDGE_LEN,
         offset,
@@ -79,7 +79,7 @@ def downsample_test_helper(use_compress):
     assert np.any(target_buffer != 0)
 
     assert np.all(
-        target_buffer == downsample_cube(source_buffer, 2, InterpolationModes.MAX)
+        target_buffer == downsample_cube(source_buffer, (2, 2, 2), InterpolationModes.MAX)
     )
 
 
@@ -114,7 +114,7 @@ def test_downsample_multi_channel():
     downsample_cube_job(
         source_info,
         target_info,
-        2,
+        (2, 2, 2),
         InterpolationModes.MAX,
         CUBE_EDGE_LEN,
         tuple(a * WKW_CUBE_SIZE for a in offset),
@@ -123,7 +123,7 @@ def test_downsample_multi_channel():
 
     channels = []
     for channel_index in range(num_channels):
-        channels.append(downsample_cube(source_data[channel_index], 2, InterpolationModes.MAX))
+        channels.append(downsample_cube(source_data[channel_index], (2, 2, 2), InterpolationModes.MAX))
     joined_buffer = np.stack(channels)
 
     target_buffer = read_wkw(

wkcuber/downsampling.py

@@ -67,10 +67,6 @@ def create_parser():
         default="uint8",
     )
 
-    parser.add_argument(
-        "--max", "-m", help="Max resolution to be downsampled", type=int, default=512
-    )
-
     parser.add_argument(
         "--from_mag",
         "--from",
@@ -80,6 +76,17 @@ def create_parser():
         default=1,
     )
 
+    # Either provide the maximum resolution to be downsampled OR a specific, anisotropic magnification.
+    group = parser.add_mutually_exclusive_group()
+    group.add_argument(
+        "--max", "-m", help="Max resolution to be downsampled", type=int, default=512
+    )
+    group.add_argument(
+        "--anisotropic_target_mag",
+        help="Specify an anisotropic target magnification which should be created (e.g., --anisotropic_target_mag 2-2-1)",
+        type=str,
+    )
+
     parser.add_argument(
         "--buffer_cube_size",
         "-b",
@@ -106,25 +113,61 @@ def cube_addresses(source_wkw_info):
         return wkw_addresses
 
 
+def mag_as_vector(mag):
+    if isinstance(mag, int):
+        return (mag, mag, mag)
+    else:
+        return mag
+
+
+def is_mag_greater_than(mag1, mag2):
+    for (m1, m2) in zip(mag1, mag2):
+        if m1 > m2:
+            return True
+    return False
+
+
+def assert_valid_mag(mag):
+    if isinstance(mag, int):
+        assert log2(mag) % 1 == 0, "magnification needs to be power of 2."
+    else:
+        assert len(mag) == 3, "magnification must be int or a vector3 of ints"
+        for mag_dim in mag:
+            assert log2(mag_dim) % 1 == 0, "magnification needs to be power of 2."
+
+
+def mag_as_layer_name(mag):
+    if isinstance(mag, int):
+        return str(mag)
+    else:
+        x, y, z = mag
+        return "{}-{}-{}".format(x, y, z)
+
+
 def downsample(
     source_wkw_info,
     target_wkw_info,
-    source_mag,
-    target_mag,
+    _source_mag,
+    _target_mag,
     interpolation_mode,
     cube_edge_len,
     jobs,
     compress,
 ):
-    assert source_mag < target_mag
-    logging.info("Downsampling mag {} from mag {}".format(target_mag, source_mag))
+    source_mag = mag_as_vector(_source_mag)
+    target_mag = mag_as_vector(_target_mag)
+    assert not is_mag_greater_than(source_mag, target_mag)
+    logging.info("Downsampling mag {} from mag {}".format(_target_mag, _source_mag))
 
-    mag_factor = int(target_mag / source_mag)
+    mag_factors = [int(t / s) for (t, s) in zip(target_mag, source_mag)]
     # Detect the cubes that we want to downsample
     source_cube_addresses = cube_addresses(source_wkw_info)
 
     target_cube_addresses = list(
-        set(tuple(x // mag_factor for x in xyz) for xyz in source_cube_addresses)
+        set(
+            tuple(dim // mag_factor for (dim, mag_factor) in zip(xyz, mag_factors))
+            for xyz in source_cube_addresses
+        )
     )
     target_cube_addresses.sort()
     logging.debug(
@@ -150,7 +193,7 @@ def downsample(
                 downsample_cube_job,
                 source_wkw_info,
                 target_wkw_info,
-                mag_factor,
+                mag_factors,
                 interpolation_mode,
                 cube_edge_len,
                 target_cube_xyz,
@@ -163,7 +206,7 @@ def downsample(
 def downsample_cube_job(
     source_wkw_info,
     target_wkw_info,
-    mag_factor,
+    mag_factors,
     interpolation_mode,
     cube_edge_len,
     target_cube_xyz,
@@ -204,12 +247,12 @@ def downsample_cube_job(
                     target_offset = np.array(
                         tile
                     ) * tile_length + wkw_cubelength * np.array(target_cube_xyz)
-                    source_offset = mag_factor * target_offset
+                    source_offset = mag_factors * target_offset
 
                     # Read source buffer
                     cube_buffer_channels = source_wkw.read(
                         source_offset,
-                        (wkw_cubelength * mag_factor // tile_count_per_dim,) * 3,
+                        (wkw_cubelength * np.array(mag_factors) // tile_count_per_dim),
                     )
 
                     for channel_index in range(num_channels):
@@ -225,7 +268,7 @@ def downsample_cube_job(
                             # Downsample the buffer
 
                             data_cube = downsample_cube(
-                                cube_buffer, mag_factor, interpolation_mode
+                                cube_buffer, mag_factors, interpolation_mode
                             )
 
                             buffer_offset = target_offset - file_offset
@@ -247,29 +290,46 @@ def downsample_cube_job(
         raise exc
 
 
-def non_linear_filter_3d(data, factor, func):
+def non_linear_filter_3d(data, factors, func):
     ds = data.shape
-    assert not any((d % factor > 0 for d in ds))
-    data = data.reshape((ds[0], factor, ds[1] // factor, ds[2]), order="F")
+    assert not any((d % factor > 0 for (d, factor) in zip(ds, factors)))
+    data = data.reshape((ds[0], factors[1], ds[1] // factors[1], ds[2]), order="F")
     data = data.swapaxes(0, 1)
     data = data.reshape(
-        (factor * factor, ds[0] * ds[1] // (factor * factor), factor, ds[2] // factor),
+        (
+            factors[0] * factors[1],
+            ds[0] * ds[1] // (factors[0] * factors[1]),
+            factors[2],
+            ds[2] // factors[2],
+        ),
         order="F",
     )
     data = data.swapaxes(2, 1)
     data = data.reshape(
         (
-            factor * factor * factor,
-            (ds[0] * ds[1] * ds[2]) // (factor * factor * factor),
+            factors[0] * factors[1] * factors[2],
+            (ds[0] * ds[1] * ds[2]) // (factors[0] * factors[1] * factors[2]),
         ),
         order="F",
     )
     data = func(data)
-    data = data.reshape((ds[0] // factor, ds[1] // factor, ds[2] // factor), order="F")
+    data = data.reshape(
+        (ds[0] // factors[0], ds[1] // factors[1], ds[2] // factors[2]), order="F"
+    )
     return data
 
 
-def linear_filter_3d(data, factor, order):
+def linear_filter_3d(data, factors, order):
+    factors = np.array(factors)
+
+    if not np.all(factors == factors[0]):
+        logging.debug(
+            "the selected filtering strategy does not support anisotropic downsampling. Selecting {} as uniform downsampling factor".format(
+                factors[0]
+            )
+        )
+    factor = factors[0]
+
     ds = data.shape
     assert not any((d % factor > 0 for d in ds))
     return zoom(
@@ -303,21 +363,21 @@ def _mode(x):
     return mode(x, axis=0, nan_policy="omit")[0][0]
 
 
-def downsample_cube(cube_buffer, factor, interpolation_mode):
+def downsample_cube(cube_buffer, factors, interpolation_mode):
     if interpolation_mode == InterpolationModes.MODE:
-        return non_linear_filter_3d(cube_buffer, factor, _mode)
+        return non_linear_filter_3d(cube_buffer, factors, _mode)
     elif interpolation_mode == InterpolationModes.MEDIAN:
-        return non_linear_filter_3d(cube_buffer, factor, _median)
+        return non_linear_filter_3d(cube_buffer, factors, _median)
     elif interpolation_mode == InterpolationModes.NEAREST:
-        return linear_filter_3d(cube_buffer, factor, 0)
+        return linear_filter_3d(cube_buffer, factors, 0)
     elif interpolation_mode == InterpolationModes.BILINEAR:
-        return linear_filter_3d(cube_buffer, factor, 1)
+        return linear_filter_3d(cube_buffer, factors, 1)
     elif interpolation_mode == InterpolationModes.BICUBIC:
-        return linear_filter_3d(cube_buffer, factor, 2)
+        return linear_filter_3d(cube_buffer, factors, 2)
     elif interpolation_mode == InterpolationModes.MAX:
-        return non_linear_filter_3d(cube_buffer, factor, _max)
+        return non_linear_filter_3d(cube_buffer, factors, _max)
     elif interpolation_mode == InterpolationModes.MIN:
-        return non_linear_filter_3d(cube_buffer, factor, _min)
+        return non_linear_filter_3d(cube_buffer, factors, _min)
     else:
         raise Exception("Invalid interpolation mode: {}".format(interpolation_mode))
 
@@ -342,8 +402,12 @@ def downsample_mag(
     else:
         interpolation_mode = InterpolationModes[interpolation_mode.upper()]
 
-    source_wkw_info = WkwDatasetInfo(path, layer_name, dtype, source_mag)
-    target_wkw_info = WkwDatasetInfo(path, layer_name, dtype, target_mag)
+    source_wkw_info = WkwDatasetInfo(
+        path, layer_name, dtype, mag_as_layer_name(source_mag)
+    )
+    target_wkw_info = WkwDatasetInfo(
+        path, layer_name, dtype, mag_as_layer_name(target_mag)
+    )
     downsample(
         source_wkw_info,
         target_wkw_info,
@@ -391,14 +455,31 @@ def downsample_mags(
     if args.verbose:
         logging.basicConfig(level=logging.DEBUG)
 
-    downsample_mags(
-        args.path,
-        args.layer_name,
-        args.from_mag,
-        args.max,
-        args.dtype,
-        args.interpolation_mode,
-        args.buffer_cube_size,
-        args.jobs,
-        args.compress,
-    )
+    if args.anisotropic_target_mag:
+        anisotropic_target_mag = tuple(map(int, args.anisotropic_target_mag.split("-")))
+        assert_valid_mag(args.from_mag)
+        assert_valid_mag(anisotropic_target_mag)
+
+        downsample_mag(
+            args.path,
+            args.layer_name,
+            args.from_mag,
+            anisotropic_target_mag,
+            args.dtype,
+            args.interpolation_mode,
+            args.buffer_cube_size,
+            args.jobs,
+            args.compress,
+        )
+    else:
+        downsample_mags(
+            args.path,
+            args.layer_name,
+            args.from_mag,
+            args.max,
+            args.dtype,
+            args.interpolation_mode,
+            args.buffer_cube_size,
+            args.jobs,
+            args.compress,
+        )