Fix bounding box of mag != 1 when exporting wkw to tiff (#205)

philippotto · web-flow · commit 545f72eaa5b8 · 2020-04-28T14:55:05.000+02:00
* fix bounding box of mag != 1 when exporting wkw to tiff

* fix formatting
diff --git a/test.sh b/test.sh
@@ -1,3 +1,4 @@
+set -euxo pipefail
 export CIRCLE_BUILD_NUM=latest
 tests/scripts/build_docker_image.sh
 tests/scripts/decompress_test_data.sh
diff --git a/tests/test_export_wkw_as_tiff.py b/tests/test_export_wkw_as_tiff.py
@@ -1,44 +1,46 @@
 from wkcuber.export_wkw_as_tiff import run, wkw_name_and_bbox_to_tiff_name
+from wkcuber.api.bounding_box import BoundingBox
 import os
 from PIL import Image
 from wkcuber.mag import Mag
 import wkw
 import numpy as np
 from math import ceil
 
+ds_name = "simple_wk_dataset"
+source_path = os.path.join("testdata", ds_name)
+
 
 def test_export_tiff_stack():
-    destination_path = os.path.join("testoutput", "WT1_wkw")
+    destination_path = os.path.join("testoutput", ds_name + "_tiff")
+    bbox = BoundingBox((100, 100, 10), (100, 500, 50))
+    bbox_dict = bbox.as_config()
     args_list = [
         "--source_path",
-        os.path.join("testdata", "WT1_wkw"),
+        source_path,
         "--destination_path",
         destination_path,
         "--layer_name",
         "color",
         "--name",
         "test_export",
         "--bbox",
-        "0,0,0,100,100,5",
+        bbox.as_csv(),
         "--mag",
         "1",
     ]
 
-    bbox = {"topleft": [0, 0, 0], "size": [100, 100, 5]}
-
     run(args_list)
 
-    test_wkw_file_path = os.path.join(
-        "testdata", "WT1_wkw", "color", Mag(1).to_layer_name()
-    )
+    test_wkw_file_path = os.path.join(source_path, "color", Mag(1).to_layer_name())
     with wkw.Dataset.open(test_wkw_file_path) as dataset:
-        slice_bbox = bbox
+        slice_bbox = bbox_dict
         slice_bbox["size"] = [slice_bbox["size"][0], slice_bbox["size"][1], 1]
-        for data_slice_index in range(1, bbox["size"][2] + 1):
+        for data_slice_index in range(1, bbox_dict["size"][2] + 1):
             slice_bbox["topleft"] = [
                 slice_bbox["topleft"][0],
                 slice_bbox["topleft"][1],
-                bbox["topleft"][2] + data_slice_index,
+                bbox_dict["topleft"][2] + data_slice_index,
             ]
             tiff_path = os.path.join(
                 destination_path,
@@ -49,8 +51,7 @@ def test_export_tiff_stack():
                 tiff_path
             ), f"Expected a tiff to be written at: {tiff_path}."
 
-            test_image = np.array(Image.open(tiff_path))
-            test_image.transpose((1, 0))
+            test_image = np.array(Image.open(tiff_path)).T
 
             correct_image = dataset.read(
                 off=slice_bbox["topleft"], shape=slice_bbox["size"]
@@ -64,10 +65,10 @@ def test_export_tiff_stack():
 
 
 def test_export_tiff_stack_tile_size():
-    destination_path = os.path.join("testoutput", "WT1_wkw_tile_size")
+    destination_path = os.path.join("testoutput", ds_name + "_tile_size")
     args_list = [
         "--source_path",
-        os.path.join("testdata", "WT1_wkw"),
+        source_path,
         "--destination_path",
         destination_path,
         "--layer_name",
@@ -87,9 +88,7 @@ def test_export_tiff_stack_tile_size():
     run(args_list)
 
     tile_bbox = {"topleft": bbox["topleft"], "size": [30, 30, 1]}
-    test_wkw_file_path = os.path.join(
-        "testdata", "WT1_wkw", "color", Mag(1).to_layer_name()
-    )
+    test_wkw_file_path = os.path.join(source_path, "color", Mag(1).to_layer_name())
     with wkw.Dataset.open(test_wkw_file_path) as dataset:
         slice_bbox = {"topleft": bbox["topleft"], "size": bbox["size"]}
         slice_bbox["size"] = [slice_bbox["size"][0], slice_bbox["size"][1], 1]
@@ -108,8 +107,7 @@ def test_export_tiff_stack_tile_size():
                         tiff_path
                     ), f"Expected a tiff to be written at: {tiff_path}."
 
-                    test_image = np.array(Image.open(tiff_path))
-                    test_image.transpose((1, 0))
+                    test_image = np.array(Image.open(tiff_path)).T
 
                     correct_image = dataset.read(
                         off=[
@@ -131,10 +129,10 @@ def test_export_tiff_stack_tile_size():
 
 
 def test_export_tiff_stack_tiles_per_dimension():
-    destination_path = os.path.join("testoutput", "WT1_wkw_tiles_per_dimension")
+    destination_path = os.path.join("testoutput", ds_name + "_tiles_per_dimension")
     args_list = [
         "--source_path",
-        os.path.join("testdata", "WT1_wkw"),
+        source_path,
         "--destination_path",
         destination_path,
         "--layer_name",
@@ -154,9 +152,7 @@ def test_export_tiff_stack_tiles_per_dimension():
     run(args_list)
 
     tile_bbox = {"topleft": bbox["topleft"], "size": [17, 17, 1]}
-    test_wkw_file_path = os.path.join(
-        "testdata", "WT1_wkw", "color", Mag(1).to_layer_name()
-    )
+    test_wkw_file_path = os.path.join(source_path, "color", Mag(1).to_layer_name())
     with wkw.Dataset.open(test_wkw_file_path) as dataset:
         slice_bbox = bbox
         slice_bbox["size"] = [slice_bbox["size"][0], slice_bbox["size"][1], 1]
@@ -177,8 +173,7 @@ def test_export_tiff_stack_tiles_per_dimension():
                         tiff_path
                     ), f"Expected a tiff to be written at: {tiff_path}."
 
-                    test_image = np.array(Image.open(tiff_path))
-                    test_image.transpose((1, 0))
+                    test_image = np.array(Image.open(tiff_path)).T
 
                     correct_image = dataset.read(
                         off=[
@@ -200,4 +195,6 @@ def test_export_tiff_stack_tiles_per_dimension():
 
 
 if __name__ == "__main__":
+    test_export_tiff_stack()
     test_export_tiff_stack_tile_size()
+    test_export_tiff_stack_tiles_per_dimension()
diff --git a/wkcuber/api/bounding_box.py b/wkcuber/api/bounding_box.py
@@ -120,6 +120,10 @@ def as_tuple6(self) -> Tuple[int, int, int, int, int, int]:
 
         return tuple(self.topleft.tolist() + self.size.tolist())  # type: ignore
 
+    def as_csv(self) -> str:
+
+        return ",".join(map(str, self.as_tuple6()))
+
     def __repr__(self) -> str:
 
         return "BoundingBox(topleft={}, size={})".format(
diff --git a/wkcuber/export_wkw_as_tiff.py b/wkcuber/export_wkw_as_tiff.py
@@ -1,5 +1,5 @@
 from argparse import ArgumentParser
-
+from functools import partial
 import logging
 import wkw
 import os
@@ -17,6 +17,7 @@
     add_batch_size_flag,
 )
 from wkcuber.mag import Mag
+from wkcuber.utils import wait_and_ensure_success
 
 
 def create_parser():
@@ -115,30 +116,30 @@ def wkw_slice_to_image(data_slice: np.ndarray, downsample: int = 1) -> Image:
 
 
 def export_tiff_slice(
-    export_args: Tuple[
-        int,
-        Tuple[
-            Dict[str, Tuple[int, int, int]],
-            str,
-            str,
-            str,
-            Union[None, Tuple[int, int]],
-            int,
-            int,
-        ],
-    ]
+    tiff_bbox: Dict[str, Tuple[int, int, int]],
+    dest_path: str,
+    name: str,
+    dataset_path: str,
+    tiling_size: Union[None, Tuple[int, int]],
+    batch_size: int,
+    downsample: int,
+    mag: Mag,
+    batch_number: int,
 ):
-    (
-        batch_number,
-        (tiff_bbox, dest_path, name, dataset_path, tiling_size, batch_size, downsample),
-    ) = export_args
+
     tiff_bbox = tiff_bbox.copy()
-    number_of_slices = min(tiff_bbox["size"][2] - batch_number * batch_size, batch_size)
-    tiff_bbox["size"] = [tiff_bbox["size"][0], tiff_bbox["size"][1], number_of_slices]
+    number_of_slices = (
+        min(tiff_bbox["size"][2] - batch_number * batch_size, batch_size) // mag.mag[2]
+    )
+    tiff_bbox["size"] = [
+        tiff_bbox["size"][0] // mag.mag[0],
+        tiff_bbox["size"][1] // mag.mag[1],
+        number_of_slices,
+    ]
     tiff_bbox["topleft"] = [
-        tiff_bbox["topleft"][0],
-        tiff_bbox["topleft"][1],
-        tiff_bbox["topleft"][2] + batch_number * batch_size,
+        tiff_bbox["topleft"][0] // mag.mag[0],
+        tiff_bbox["topleft"][1] // mag.mag[1],
+        (tiff_bbox["topleft"][2] + batch_number * batch_size) // mag.mag[2],
     ]
 
     with wkw.Dataset.open(dataset_path) as dataset:
@@ -155,7 +156,6 @@ def export_tiff_slice(
             slice_name_number = batch_number * batch_size + slice_index + 1
             if tiling_size is None:
                 tiff_file_name = wkw_name_and_bbox_to_tiff_name(name, slice_name_number)
-
                 tiff_file_path = os.path.join(dest_path, tiff_file_name)
 
                 image = wkw_slice_to_image(tiff_data[:, :, :, slice_index], downsample)
@@ -213,23 +213,23 @@ def export_tiff_stack(
     with get_executor_for_args(args) as executor:
         num_slices = ceil(bbox["size"][2] / batch_size)
         slices = range(0, num_slices)
-        export_args = zip(
+
+        logging.info(f"starting jobs")
+        futures = executor.map_to_futures(
+            partial(
+                export_tiff_slice,
+                bbox,
+                destination_path,
+                name,
+                dataset_path,
+                tiling_slice_size,
+                batch_size,
+                downsample,
+                mag,
+            ),
             slices,
-            [
-                (
-                    bbox,
-                    destination_path,
-                    name,
-                    dataset_path,
-                    tiling_slice_size,
-                    batch_size,
-                    downsample,
-                )
-            ]
-            * num_slices,
         )
-        logging.info(f"starting jobs")
-        executor.map(export_tiff_slice, export_args)
+        wait_and_ensure_success(futures)
 
 
 def export_wkw_as_tiff(args):

Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,4 @@`
	`1`	`+set -euxo pipefail`
`1`	`2`	`export CIRCLE_BUILD_NUM=latest`
`2`	`3`	`tests/scripts/build_docker_image.sh`
`3`	`4`	`tests/scripts/decompress_test_data.sh`