Merge pull request #55 from scalableminds/mag

Mag

Author: valentin-pinkau

setup.py

@@ -1,18 +1,12 @@
 from setuptools import setup, find_packages
 
 setup(
-    name='wkcuber',
+    name="wkcuber",
     packages=find_packages(exclude=("tests",)),
-    version='0.2.0',
-    install_requires=[
-        'scipy',
-        'numpy',
-        'pillow',
-        'pyyaml',
-        'wkw'
-    ],
-    description='A cubing tool for webKnossos',
-    author='Norman Rzepka',
-    author_email='[email protected]',
-    url='https://scalableminds.com'
+    version="0.2.0",
+    install_requires=["scipy", "numpy", "pillow", "pyyaml", "wkw"],
+    description="A cubing tool for webKnossos",
+    author="Norman Rzepka",
+    author_email="[email protected]",
+    url="https://scalableminds.com",
 )

testdata/WT1_wkw.tar.gz

@@ -79,7 +79,8 @@ def downsample_test_helper(use_compress):
     assert np.any(target_buffer != 0)
 
     assert np.all(
-        target_buffer == downsample_cube(source_buffer, (2, 2, 2), InterpolationModes.MAX)
+        target_buffer
+        == downsample_cube(source_buffer, (2, 2, 2), InterpolationModes.MAX)
     )
 
 
@@ -103,10 +104,14 @@ def test_downsample_multi_channel():
     offset = (0, 0, 0)
     num_channels = 3
     size = (32, 32, 10)
-    source_data = (128 * np.random.randn(num_channels, size[0], size[1], size[2])).astype('uint8')
+    source_data = (
+        128 * np.random.randn(num_channels, size[0], size[1], size[2])
+    ).astype("uint8")
     file_len = 32
 
-    with open_wkw(source_info, num_channels=num_channels, file_len=file_len) as wkw_dataset:
+    with open_wkw(
+        source_info, num_channels=num_channels, file_len=file_len
+    ) as wkw_dataset:
         print("writing source_data shape", source_data.shape)
         wkw_dataset.write(offset, source_data)
     assert np.any(source_data != 0)
@@ -123,17 +128,19 @@ def test_downsample_multi_channel():
 
     channels = []
     for channel_index in range(num_channels):
-        channels.append(downsample_cube(source_data[channel_index], (2, 2, 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(
         target_info,
         tuple(a * WKW_CUBE_SIZE for a in offset),
         list(map(lambda x: x // 2, size)),
-        file_len=file_len
+        file_len=file_len,
     )
     assert np.any(target_buffer != 0)
 
-    assert np.all(
-        target_buffer == joined_buffer
-    )
+    assert np.all(target_buffer == joined_buffer)

tests/test_downsampling.py

@@ -0,0 +1,25 @@
+from wkcuber.mag import Mag
+from wkcuber.metadata import detect_resolutions
+
+
+def test_detect_resolutions():
+    resolutions = sorted(list(detect_resolutions("/testdata/WT1_wkw", "color")))
+    assert [mag.to_layer_name() for mag in resolutions] == ["1", "2-2-1"]
+
+
+def test_mag_constructor():
+    mag = Mag(16)
+    assert mag.to_array() == [16, 16, 16]
+
+    mag = Mag("256")
+    assert mag.to_array() == [256, 256, 256]
+
+    mag = Mag("16-2-4")
+
+    assert mag.to_array() == [16, 2, 4]
+
+    mag1 = Mag("16-2-4")
+    mag2 = Mag("8-2-4")
+
+    assert mag1 > mag2
+    assert mag1.to_layer_name() == "16-2-4"

tests/test_mag.py

@@ -1,14 +1,13 @@
 from argparse import ArgumentParser
-from os import path
-from uuid import uuid4
+
 import logging
-import shutil
 
 from .cubing import cubing, BLOCK_LEN
 from .downsampling import downsample_mags, DEFAULT_EDGE_LEN
-from .compress import compress_mag
+from .compress import compress_mag_inplace
 from .metadata import write_webknossos_metadata
 from .utils import add_verbose_flag, add_jobs_flag
+from .mag import Mag
 
 
 def create_parser():
@@ -72,18 +71,6 @@ def create_parser():
     return parser
 
 
-def compress_mag_inplace(target_path, layer_name, mag, jobs):
-    compress_target_path = "{}.compress-{}".format(target_path, uuid4())
-    compress_mag(target_path, layer_name, compress_target_path, mag, jobs)
-
-    shutil.rmtree(path.join(args.target_path, args.layer_name, str(mag)))
-    shutil.move(
-        path.join(compress_target_path, layer_name, str(mag)),
-        path.join(target_path, layer_name, str(mag)),
-    )
-    shutil.rmtree(compress_target_path)
-
-
 if __name__ == "__main__":
     args = create_parser().parse_args()
 
@@ -100,13 +87,13 @@ def compress_mag_inplace(target_path, layer_name, mag, jobs):
     )
 
     if not args.no_compress:
-        compress_mag_inplace(args.target_path, args.layer_name, 1, args.jobs)
+        compress_mag_inplace(args.target_path, args.layer_name, Mag(1), args.jobs)
 
     downsample_mags(
         args.target_path,
         args.layer_name,
-        1,
-        args.max_mag,
+        Mag(1),
+        Mag(args.max_mag),
         args.dtype,
         "default",
         DEFAULT_EDGE_LEN,

wkcuber/__main__.py

@@ -5,6 +5,8 @@
 import numpy as np
 from argparse import ArgumentParser
 from os import path, makedirs
+from uuid import uuid4
+from .mag import Mag
 
 from .utils import (
     add_verbose_flag,
@@ -14,6 +16,7 @@
     ParallelExecutor,
 )
 from .metadata import detect_resolutions
+from typing import List
 
 
 def create_parser():
@@ -68,25 +71,39 @@ def compress_file_job(source_path, target_path):
         raise exc
 
 
-def compress_mag(source_path, layer_name, target_path, mag, jobs):
+def compress_mag(source_path, layer_name, target_path, mag: Mag, jobs):
     if path.exists(path.join(target_path, layer_name, str(mag))):
         logging.error("Target path '{}' already exists".format(target_path))
         exit(1)
 
     source_wkw_info = WkwDatasetInfo(source_path, layer_name, None, mag)
     target_mag_path = path.join(target_path, layer_name, str(mag))
-    logging.info("Compressing mag {0} in '{1}'".format(mag, target_mag_path))
+    logging.info("Compressing mag {0} in '{1}'".format(str(mag), target_mag_path))
 
     with open_wkw(source_wkw_info) as source_wkw, ParallelExecutor(jobs) as pool:
         source_wkw.compress(target_mag_path)
         for file in source_wkw.list_files():
             rel_file = path.relpath(file, source_wkw.root)
             pool.submit(compress_file_job, file, path.join(target_mag_path, rel_file))
 
-    logging.info("Mag {0} succesfully compressed".format(mag))
+    logging.info("Mag {0} succesfully compressed".format(str(mag)))
 
 
-def compress_mags(source_path, layer_name, target_path=None, mags=None, jobs=1):
+def compress_mag_inplace(target_path, layer_name, mag: Mag, jobs):
+    compress_target_path = "{}.compress-{}".format(target_path, uuid4())
+    compress_mag(target_path, layer_name, compress_target_path, mag, jobs)
+
+    shutil.rmtree(path.join(target_path, layer_name, str(mag)))
+    shutil.move(
+        path.join(compress_target_path, layer_name, str(mag)),
+        path.join(target_path, layer_name, str(mag)),
+    )
+    shutil.rmtree(compress_target_path)
+
+
+def compress_mags(
+    source_path, layer_name, target_path=None, mags: List[Mag] = None, jobs=1
+):
     with_tmp_dir = target_path is None
     target_path = source_path + ".tmp" if with_tmp_dir else target_path
 

wkcuber/compress.py

@@ -10,6 +10,7 @@
 from itertools import product
 from functools import lru_cache
 from enum import Enum
+from .mag import Mag
 
 from .utils import (
     add_jobs_flag,
@@ -113,53 +114,23 @@ 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: Mag,
+    target_mag: Mag,
     interpolation_mode,
     cube_edge_len,
     jobs,
     compress,
 ):
-    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_factors = [int(t / s) for (t, s) in zip(target_mag, source_mag)]
+    assert source_mag < target_mag
+    logging.info("Downsampling mag {} from mag {}".format(target_mag, source_mag))
+
+    mag_factors = [
+        t // s for (t, s) in zip(target_mag.to_array(), source_mag.to_array())
+    ]
     # Detect the cubes that we want to downsample
     source_cube_addresses = cube_addresses(source_wkw_info)
 
@@ -385,8 +356,8 @@ def downsample_cube(cube_buffer, factors, interpolation_mode):
 def downsample_mag(
     path,
     layer_name,
-    source_mag,
-    target_mag,
+    source_mag: Mag,
+    target_mag: Mag,
     dtype="uint8",
     interpolation_mode="default",
     cube_edge_len=DEFAULT_EDGE_LEN,
@@ -403,10 +374,10 @@ def downsample_mag(
         interpolation_mode = InterpolationModes[interpolation_mode.upper()]
 
     source_wkw_info = WkwDatasetInfo(
-        path, layer_name, dtype, mag_as_layer_name(source_mag)
+        path, layer_name, dtype, source_mag.to_layer_name()
     )
     target_wkw_info = WkwDatasetInfo(
-        path, layer_name, dtype, mag_as_layer_name(target_mag)
+        path, layer_name, dtype, target_mag.to_layer_name()
     )
     downsample(
         source_wkw_info,
@@ -423,18 +394,17 @@ def downsample_mag(
 def downsample_mags(
     path,
     layer_name,
-    from_mag,
-    max_mag,
+    from_mag: Mag,
+    max_mag: Mag,
     dtype,
     interpolation_mode,
     cube_edge_len,
     jobs,
     compress,
 ):
-    assert log2(from_mag) % 1 == 0, "'from_mag' needs to be power of 2."
-    target_mag = from_mag * 2
+    target_mag = from_mag.scaled_by(2)
     while target_mag <= max_mag:
-        source_mag = target_mag // 2
+        source_mag = target_mag.divided_by(2)
         downsample_mag(
             path,
             layer_name,
@@ -446,7 +416,7 @@ def downsample_mags(
             jobs,
             compress,
         )
-        target_mag = target_mag * 2
+        target_mag.scale_by(2)
 
 
 if __name__ == "__main__":
@@ -455,15 +425,15 @@ def downsample_mags(
     if args.verbose:
         logging.basicConfig(level=logging.DEBUG)
 
+    from_mag = Mag(args.from_mag)
+    max_mag = Mag(args.max)
     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)
+        anisotropic_target_mag = Mag(args.anisotropic_target_mag)
 
         downsample_mag(
             args.path,
             args.layer_name,
-            args.from_mag,
+            from_mag,
             anisotropic_target_mag,
             args.dtype,
             args.interpolation_mode,
@@ -475,8 +445,8 @@ def downsample_mags(
         downsample_mags(
             args.path,
             args.layer_name,
-            args.from_mag,
-            args.max,
+            from_mag,
+            max_mag,
             args.dtype,
             args.interpolation_mode,
             args.buffer_cube_size,