Add pad option to cube images with differing size (#94)

* add pad option to cube images with differing size

* refactor argparse and fix missing arg for main

* fix dimensions in padding

* bump version

Author: philippotto

setup.py

@@ -3,7 +3,7 @@
 setup(
     name="wkcuber",
     packages=find_packages(exclude=("tests",)),
-    version="0.2.12",
+    version="0.2.13",
     install_requires=["scipy", "numpy", "pillow", "pyyaml", "wkw", "cluster_tools>=1.19"],
     description="A cubing tool for webKnossos",
     author="Norman Rzepka",

wkcuber/__main__.py

@@ -2,7 +2,7 @@
 
 import logging
 
-from .cubing import cubing, BLOCK_LEN
+from .cubing import cubing, BLOCK_LEN, create_parser as create_cubing_parser
 from .downsampling import downsample_mags, downsample_mags_anisotropic, DEFAULT_EDGE_LEN
 from .compress import compress_mag_inplace
 from .metadata import write_webknossos_metadata
@@ -11,35 +11,7 @@
 
 
 def create_parser():
-    parser = ArgumentParser()
-
-    parser.add_argument("source_path", help="Directory containing the input images.")
-
-    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(
-        "--batch_size",
-        "-b",
-        type=int,
-        help="Number of slices to buffer per job",
-        default=BLOCK_LEN,
-    )
+    parser = create_cubing_parser()
 
     parser.add_argument(
         "--max_mag",
@@ -65,8 +37,6 @@ def create_parser():
         default="1,1,1",
     )
 
-    add_verbose_flag(parser)
-    add_distribution_flags(parser)
     add_anisotropic_flag(parser)
 
     return parser

wkcuber/cubing.py

@@ -52,6 +52,15 @@ def create_parser():
         default=BLOCK_LEN,
     )
 
+    parser.add_argument(
+        "--pad",
+        help="Automatically pad image files at the bottom and right borders. "
+        "Use this, when the input images don't have a common size, but have "
+        "their origin at (0, 0).",
+        default=False,
+        action="store_true",
+    )
+
     add_verbose_flag(parser)
     add_distribution_flags(parser)
 
@@ -89,6 +98,7 @@ def cubing_job(
     batch_size,
     image_size,
     num_channels,
+    pad=False,
 ):
     if len(z_batches) == 0:
         return
@@ -109,13 +119,32 @@ def cubing_job(
                 for z, file_name in zip(z_batch, source_file_batch):
                     # Image shape will be (x, y, channel_count, z=1) or (x, y, z=1)
                     image = read_image_file(file_name, target_wkw_info.dtype)
-                    assert (
-                        image.shape[0:2] == image_size
-                    ), "Section z={} has the wrong dimensions: {} (expected {}).".format(
-                        z, image.shape, image_size
-                    )
+                    if not pad:
+                        assert (
+                            image.shape[0:2] == image_size
+                        ), "Section z={} has the wrong dimensions: {} (expected {}). Consider using --pad.".format(
+                            z, image.shape, image_size
+                        )
                     buffer.append(image)
 
+                if pad:
+                    x_max = max(slice.shape[0] for slice in buffer)
+                    y_max = max(slice.shape[1] for slice in buffer)
+
+                    buffer = [
+                        np.pad(
+                            slice,
+                            mode="constant",
+                            pad_width=[
+                                (0, x_max - slice.shape[0]),
+                                (0, y_max - slice.shape[1]),
+                                (0, 0),
+                                (0, 0),
+                            ],
+                        )
+                        for slice in buffer
+                    ]
+
                 # Write batch buffer which will have shape (x, y, channel_count, z)
                 # since we concat along the last axis (z)
                 buffer = np.concatenate(buffer, axis=-1)
@@ -175,6 +204,7 @@ def cubing(source_path, target_path, layer_name, dtype, batch_size, args=None) -
                     batch_size,
                     (num_x, num_y),
                     num_channels,
+                    args.pad,
                 )
             )