add recubing script (#140)

* add recubing script
* Merge branch 'master' into recubing
* add wkw headers in script to accomodate changed open_wkw function, reformatting
* move methods cube_addresses, parse_cube_file_name to utils, use open_wkw in with-context
* remove convert_nifti file
* remove infer_bounding_box, instead usage of detect_bbox
* add more reformatting
* add multiprocessing to recubing script
* Reformatted
* Merge branch 'master' into recubing

Author: leowe

wkcuber/downsampling.py

@@ -23,15 +23,11 @@
     wait_and_ensure_success,
     add_isotropic_flag,
     setup_logging,
+    cube_addresses,
+    parse_cube_file_name,
 )
 
 DEFAULT_EDGE_LEN = 256
-CUBE_REGEX = re.compile(r"z(\d+)/y(\d+)/x(\d+)(\.wkw)$")
-
-
-def parse_cube_file_name(filename):
-    m = CUBE_REGEX.search(filename)
-    return (int(m.group(3)), int(m.group(2)), int(m.group(1)))
 
 
 def determine_buffer_edge_len(dataset):
@@ -127,14 +123,6 @@ def create_parser():
     return parser
 
 
-def cube_addresses(source_wkw_info):
-    # Gathers all WKW cubes in the dataset
-    with open_wkw(source_wkw_info) as source_wkw:
-        wkw_addresses = list(parse_cube_file_name(f) for f in source_wkw.list_files())
-        wkw_addresses.sort()
-        return wkw_addresses
-
-
 def downsample(
     source_wkw_info,
     target_wkw_info,

wkcuber/recubing.py

@@ -0,0 +1,179 @@
+import logging
+import wkw
+import numpy as np
+from argparse import ArgumentParser
+from itertools import product
+
+from .metadata import detect_bbox
+
+from .utils import (
+    add_verbose_flag,
+    open_wkw,
+    WkwDatasetInfo,
+    ensure_wkw,
+    add_distribution_flags,
+    setup_logging,
+    get_executor_for_args,
+    wait_and_ensure_success,
+)
+
+
+def create_parser():
+    parser = ArgumentParser()
+
+    parser.add_argument(
+        "source_path", help="Directory containing the datasource properties."
+    )
+
+    parser.add_argument(
+        "target_path", help="Output directory for the generated dataset."
+    )
+
+    parser.add_argument(
+        "--layer_name",
+        "-l",
+        help="Name of the cubed layer (color or segmentation)",
+        default="color",
+    )
+
+    parser.add_argument(
+        "--dtype",
+        "-d",
+        help="Target datatype (e.g. uint8, uint16, uint32)",
+        default="uint8",
+    )
+
+    parser.add_argument(
+        "--wkw_file_len", help="Target file length (default 32)", type=int, default=32
+    )
+
+    parser.add_argument(
+        "--no_compression",
+        help="Use compression, default false",
+        type=bool,
+        default=False,
+    )
+
+    add_verbose_flag(parser)
+    add_distribution_flags(parser)
+
+    return parser
+
+
+def next_lower_divisible_by(number, divisor) -> int:
+    remainder = number % divisor
+    return number - remainder
+
+
+def next_higher_divisible_by(number, divisor) -> int:
+    remainder = number % divisor
+    return number - remainder + divisor
+
+
+def recube(
+    source_path, target_path, layer_name, dtype, wkw_file_len=32, compression=True
+):
+    if compression:
+        block_type = wkw.Header.BLOCK_TYPE_LZ4
+    else:
+        block_type = wkw.Header.BLOCK_TYPE_RAW
+
+    target_wkw_header = wkw.Header(
+        np.dtype(dtype), file_len=wkw_file_len, block_type=block_type
+    )
+    target_wkw_info = WkwDatasetInfo(target_path, layer_name, 1, target_wkw_header)
+    source_wkw_header = wkw.Header(np.dtype(dtype))
+    source_wkw_info = WkwDatasetInfo(source_path, layer_name, 1, source_wkw_header)
+
+    ensure_wkw(target_wkw_info)
+
+    bounding_box_dict = detect_bbox(source_wkw_info.dataset_path, layer_name)
+    bounding_box = (
+        bounding_box_dict["topLeft"],
+        [
+            bounding_box_dict["width"],
+            bounding_box_dict["height"],
+            bounding_box_dict["depth"],
+        ],
+    )
+    bottom_right = [
+        coord + size for coord, size in zip(bounding_box[0], bounding_box[1])
+    ]
+
+    wkw_cube_size = wkw_file_len * target_wkw_header.block_len
+
+    outer_bounding_box_tl = list(
+        map(lambda lx: next_lower_divisible_by(lx, wkw_cube_size), bounding_box[0])
+    )
+    outer_bounding_box_br = list(
+        map(lambda lx: next_higher_divisible_by(lx, wkw_cube_size), bottom_right)
+    )
+    outer_bounding_box_size = [
+        outer_bounding_box_br[0] - outer_bounding_box_tl[0],
+        outer_bounding_box_br[1] - outer_bounding_box_tl[1],
+        outer_bounding_box_br[2] - outer_bounding_box_tl[2],
+    ]
+
+    target_cube_addresses = product(
+        range(0, outer_bounding_box_size[0], wkw_cube_size),
+        range(0, outer_bounding_box_size[1], wkw_cube_size),
+        range(0, outer_bounding_box_size[2], wkw_cube_size),
+    )
+
+    with get_executor_for_args(args) as executor:
+        job_args = []
+        for target_cube_xyz in target_cube_addresses:
+            job_args.append(
+                (
+                    source_wkw_info,
+                    target_wkw_info,
+                    outer_bounding_box_size,
+                    outer_bounding_box_tl,
+                    wkw_cube_size,
+                    target_cube_xyz,
+                )
+            )
+        wait_and_ensure_success(executor.map_to_futures(recubing_cube_job, job_args))
+
+    logging.info(f"{layer_name} successfully resampled!")
+
+
+def recubing_cube_job(args):
+    (
+        source_wkw_info,
+        target_wkw_info,
+        outer_bounding_box_size,
+        outer_bounding_box_tl,
+        wkw_cube_size,
+        target_cube_xyz,
+    ) = args
+
+    with open_wkw(source_wkw_info) as source_wkw_dataset:
+        with open_wkw(target_wkw_info) as target_wkw_dataset:
+            top_left = [
+                outer_bounding_box_tl[0] + target_cube_xyz[0],
+                outer_bounding_box_tl[1] + target_cube_xyz[1],
+                outer_bounding_box_tl[2] + target_cube_xyz[2],
+            ]
+
+            logging.info("Writing at {}".format(top_left))
+
+            data_cube = source_wkw_dataset.read(
+                top_left, (wkw_cube_size, wkw_cube_size, wkw_cube_size)
+            )
+
+            target_wkw_dataset.write(top_left, data_cube)
+
+
+if __name__ == "__main__":
+    args = create_parser().parse_args()
+    setup_logging(args)
+
+    recube(
+        args.source_path,
+        args.target_path,
+        args.layer_name,
+        args.dtype,
+        args.wkw_file_len,
+        not args.no_compression,
+    )

wkcuber/utils.py

@@ -1,3 +1,4 @@
+import re
 import time
 import wkw
 import numpy as np
@@ -29,6 +30,7 @@
 
 
 BLOCK_LEN = 32
+CUBE_REGEX = re.compile(r"z(\d+)/y(\d+)/x(\d+)(\.wkw)$")
 
 logger = getLogger(__name__)
 
@@ -47,6 +49,19 @@ def ensure_wkw(target_wkw_info):
     target_wkw.close()
 
 
+def cube_addresses(source_wkw_info):
+    # Gathers all WKW cubes in the dataset
+    with open_wkw(source_wkw_info) as source_wkw:
+        wkw_addresses = list(parse_cube_file_name(f) for f in source_wkw.list_files())
+        wkw_addresses.sort()
+        return wkw_addresses
+
+
+def parse_cube_file_name(filename):
+    m = CUBE_REGEX.search(filename)
+    return int(m.group(3)), int(m.group(2)), int(m.group(1))
+
+
 def open_knossos(info):
     return KnossosDataset.open(info.dataset_path, np.dtype(info.dtype))