From cd69fb754f8047f395ec2cc04b094032c4f27ef2 Mon Sep 17 00:00:00 2001 From: R Schwanhold Date: Thu, 5 Mar 2020 10:54:09 +0100 Subject: [PATCH] High level dataset api chunking (#184) * only log downsampling for each MB ob precessed data * log every gigabyte instead of every megabyte * implement PR feedback * WIP: implement high level dataset api * WIP: implement TiffMag * implement test for dataset api * refactor high level dataset api * add test data for WKDataset * improve high level dataset api and add more tests * use seed in tests to achieve deterministic behaviour * implement a test for writing out of bounds with a wk_slice * add possibility to define pattern for TiffDatasets * implement tiled tiff * improve quality of dataset tests * make naming schema of test files consistent * implement PR feedback * rename tiffs in testdata * Improve error message and reformat code * add type annotation * rename Slice to View * remove comments and reformat code * add scikit-image (which is needed for skimage) as dependency * update version of scikit-image * fix relative paths in tests * add possibility to get data of specific tile * reformat code * add test for png images * reformat code * make TiledTiffDataset a seperate class * add support for opening datasets which do not have num_channels in properties.json * add test for advanced pattern * use more efficient method to detect tile_ranges * rename file path of test case * support largestSegmentationId in properties * restructure properties * add dataType attribute to properties * WIP: implement chunking * implement PR feedback * remove comment * fix relative paths * fix writing data to an unopened dataset and support different orders of the dimensions with patterns * Adjust docstring of method Co-Authored-By: Philipp Otto * use type to call subclass constructor * reformat code * Support tiled tiffs in high level dataset API (#178) * add possibility to define pattern for TiffDatasets * implement tiled tiff * add possibility to get data of specific tile * add test for png images * reformat code * make TiledTiffDataset a seperate class * add support for opening datasets which do not have num_channels in properties.json * add test for advanced pattern * use more efficient method to detect tile_ranges * rename file path of test case * support largestSegmentationId in properties * restructure properties * add dataType attribute to properties * implement PR feedback * fix writing data to an unopened dataset and support different orders of the dimensions with patterns * Adjust docstring of method Co-Authored-By: Philipp Otto Co-authored-by: Philipp Otto * implement get_view for a View; fix check_chunk_size for WKView * reformat code Co-authored-by: Philipp Otto --- tests/test_dataset.py | 251 ++++++++++++++++++++++++++++++++++-- wkcuber/api/Dataset.py | 7 +- wkcuber/api/MagDataset.py | 53 ++++++-- wkcuber/api/View.py | 79 ++++++++++++ wkcuber/api/bounding_box.py | 214 ++++++++++++++++++++++++++++++ 5 files changed, 579 insertions(+), 25 deletions(-) create mode 100644 wkcuber/api/bounding_box.py diff --git a/tests/test_dataset.py b/tests/test_dataset.py index e9194b85e..1a134e20f 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -11,7 +11,9 @@ from wkcuber.api.Layer import Layer from wkcuber.api.Properties.DatasetProperties import TiffProperties, WKProperties from wkcuber.api.TiffData.TiffMag import TiffReader +from wkcuber.api.bounding_box import BoundingBox from wkcuber.mag import Mag +from wkcuber.utils import get_executor_for_args def delete_dir(relative_path): @@ -19,6 +21,99 @@ def delete_dir(relative_path): rmtree(relative_path) +def chunk_job(args): + view, additional_args = args + + # increment the color value of each voxel + data = view.read(view.size) + if data.shape[0] == 1: + data = data[0, :, :, :] + data += 50 + view.write(data) + + +def advanced_chunk_job(args): + view, additional_args = args + + # write different data for each chunk (depending on the global_offset of the chunk) + data = view.read(view.size) + data = np.ones(data.shape, dtype=np.uint8) * np.uint8(sum(view.global_offset)) + view.write(data) + + +def for_each_chunking_with_wrong_chunk_size(view): + with get_executor_for_args(None) as executor: + try: + view.for_each_chunk( + chunk_job, + job_args_per_chunk="test", + chunk_size=(0, 64, 64), + executor=executor, + ) + raise Exception( + "The test did not throw an exception even though it should. " + "The chunk_size should not contain zeros" + ) + except AssertionError: + pass + + try: + view.for_each_chunk( + chunk_job, + job_args_per_chunk="test", + chunk_size=(16, 64, 64), + executor=executor, + ) + raise Exception( + "The test did not throw an exception even though it should. " + ) + except AssertionError: + pass + + try: + view.for_each_chunk( + chunk_job, + job_args_per_chunk="test", + chunk_size=(100, 64, 64), + executor=executor, + ) + raise Exception( + "The test did not throw an exception even though it should. " + ) + except AssertionError: + pass + + +def for_each_chunking_advanced(ds, view): + chunk_size = (64, 64, 64) + with get_executor_for_args(None) as executor: + view.for_each_chunk( + advanced_chunk_job, + job_args_per_chunk="test", + chunk_size=chunk_size, + executor=executor, + ) + + for offset, size in [ + ((10, 10, 10), (54, 54, 54)), + ((10, 64, 10), (54, 64, 54)), + ((10, 128, 10), (54, 32, 54)), + ((64, 10, 10), (64, 54, 54)), + ((64, 64, 10), (64, 64, 54)), + ((64, 128, 10), (64, 32, 54)), + ((128, 10, 10), (32, 54, 54)), + ((128, 64, 10), (32, 64, 54)), + ((128, 128, 10), (32, 32, 54)), + ]: + chunk = ds.get_view("color", "1", size=size, offset=offset, is_bounded=False) + chunk_data = chunk.read(chunk.size) + assert np.array_equal( + np.ones(chunk_data.shape, dtype=np.uint8) + * np.uint8(sum(chunk.global_offset)), + chunk_data, + ) + + def get_multichanneled_data(dtype): data = np.zeros((3, 250, 200, 10), dtype=dtype) for h in range(10): @@ -108,7 +203,7 @@ def test_view_read_with_open(): # This test would be the same for TiffDataset wk_view = WKDataset("./testdata/simple_wk_dataset/").get_view( - "color", "1", size=(32, 32, 32) + "color", "1", size=(16, 16, 16) ) assert not wk_view._is_opened @@ -126,7 +221,7 @@ def test_view_read_without_open(): # This test would be the same for TiffDataset wk_view = WKDataset("./testdata/simple_wk_dataset/").get_view( - "color", "1", size=(32, 32, 32) + "color", "1", size=(16, 16, 16) ) assert not wk_view._is_opened @@ -143,7 +238,7 @@ def test_view_wk_write(): copytree("./testdata/simple_wk_dataset/", "./testoutput/simple_wk_dataset/") wk_view = WKDataset("./testoutput/simple_wk_dataset/").get_view( - "color", "1", size=(100, 100, 100) + "color", "1", size=(16, 16, 16) ) with wk_view.open(): @@ -161,7 +256,7 @@ def test_view_tiff_write(): copytree("./testdata/simple_tiff_dataset/", "./testoutput/simple_tiff_dataset/") tiff_view = TiffDataset("./testoutput/simple_tiff_dataset/").get_view( - "color", "1", size=(100, 100, 100) + "color", "1", size=(16, 16, 10) ) with tiff_view.open(): @@ -182,7 +277,7 @@ def test_view_tiff_write_out_of_bounds(): copytree("./testdata/simple_tiff_dataset/", new_dataset_path) tiff_view = TiffDataset(new_dataset_path).get_view( - "color", "1", size=(100, 100, 100) + "color", "1", size=(100, 100, 10) ) with tiff_view.open(): @@ -203,7 +298,7 @@ def test_view_wk_write_out_of_bounds(): delete_dir(new_dataset_path) copytree("./testdata/simple_wk_dataset/", new_dataset_path) - tiff_view = WKDataset(new_dataset_path).get_view("color", "1", size=(100, 100, 100)) + tiff_view = WKDataset(new_dataset_path).get_view("color", "1", size=(16, 16, 16)) with tiff_view.open(): try: @@ -217,6 +312,32 @@ def test_view_wk_write_out_of_bounds(): pass +def test_wk_view_out_of_bounds(): + try: + # The size of the mag is (24, 24, 24). Trying to get an bigger view should throw an error + WKDataset("./testdata/simple_wk_dataset/").get_view( + "color", "1", size=(100, 100, 100) + ) + raise Exception( + "The test 'test_view_wk_write_out_of_bounds' did not throw an exception even though it should" + ) + except AssertionError: + pass + + +def test_tiff_view_out_of_bounds(): + try: + # The size of the mag is (24, 24, 24). Trying to get an bigger view should throw an error + TiffDataset("./testdata/simple_tiff_dataset/").get_view( + "color", "1", size=(100, 100, 100) + ) + raise Exception( + "The test 'test_view_wk_write_out_of_bounds' did not throw an exception even though it should" + ) + except AssertionError: + pass + + def test_tiff_write_out_of_bounds(): new_dataset_path = "./testoutput/simple_tiff_dataset_out_of_bounds/" @@ -690,6 +811,120 @@ def test_properties_with_segmentation(): assert input_data == output_data +def test_chunking_wk(): + delete_dir("./testoutput/chunking_dataset_wk/") + copytree("./testdata/simple_wk_dataset/", "./testoutput/chunking_dataset_wk/") + + view = WKDataset("./testoutput/chunking_dataset_wk/").get_view( + "color", "1", size=(256, 256, 256), is_bounded=False + ) + + original_data = view.read(view.size) + + with get_executor_for_args(None) as executor: + view.for_each_chunk( + chunk_job, + job_args_per_chunk="test", + chunk_size=(64, 64, 64), + executor=executor, + ) + + assert np.array_equal(original_data + 50, view.read(view.size)) + + +def test_chunking_wk_advanced(): + delete_dir("./testoutput/chunking_dataset_wk_advanced/") + copytree( + "./testdata/simple_wk_dataset/", "./testoutput/chunking_dataset_wk_advanced/" + ) + + ds = WKDataset("./testoutput/chunking_dataset_wk_advanced/") + view = ds.get_view( + "color", "1", size=(150, 150, 54), offset=(10, 10, 10), is_bounded=False + ) + for_each_chunking_advanced(ds, view) + + +def test_chunking_wk_wrong_chunk_size(): + delete_dir("./testoutput/chunking_dataset_wk_with_wrong_chunk_size/") + copytree( + "./testdata/simple_wk_dataset/", + "./testoutput/chunking_dataset_wk_with_wrong_chunk_size/", + ) + + view = WKDataset( + "./testoutput/chunking_dataset_wk_with_wrong_chunk_size/" + ).get_view("color", "1", size=(256, 256, 256), is_bounded=False) + + for_each_chunking_with_wrong_chunk_size(view) + + +def test_chunking_tiff(): + delete_dir("./testoutput/chunking_dataset_tiff/") + copytree("./testdata/simple_tiff_dataset/", "./testoutput/chunking_dataset_tiff/") + + view = TiffDataset("./testoutput/chunking_dataset_tiff/").get_view( + "color", "1", size=(265, 265, 10) + ) + + original_data = view.read(view.size) + + with get_executor_for_args(None) as executor: + view.for_each_chunk( + chunk_job, + job_args_per_chunk="test", + chunk_size=(265, 265, 1), + executor=executor, + ) + + new_data = view.read(view.size) + assert np.array_equal(original_data + 50, new_data) + + +def test_chunking_tiff_wrong_chunk_size(): + delete_dir("./testoutput/chunking_dataset_tiff_with_wrong_chunk_size/") + copytree( + "./testdata/simple_tiff_dataset/", + "./testoutput/chunking_dataset_tiff_with_wrong_chunk_size/", + ) + + view = TiffDataset( + "./testoutput/chunking_dataset_tiff_with_wrong_chunk_size/" + ).get_view("color", "1", size=(256, 256, 256), is_bounded=False) + + for_each_chunking_with_wrong_chunk_size(view) + + +def test_chunking_tiled_tiff_wrong_chunk_size(): + delete_dir("./testoutput/chunking_dataset_tiled_tiff_with_wrong_chunk_size/") + + ds = TiledTiffDataset.create( + "./testoutput/chunking_dataset_tiled_tiff_with_wrong_chunk_size/", + scale=(1, 1, 1), + tile_size=(32, 32), + pattern="{xxxx}/{yyyy}/{zzzz}.tif", + ) + ds.add_layer("color", Layer.COLOR_TYPE).add_mag("1") + view = ds.get_view("color", "1", size=(256, 256, 256), is_bounded=False) + + for_each_chunking_with_wrong_chunk_size(view) + + +def test_chunking_tiled_tiff_advanced(): + delete_dir("./testoutput/chunking_dataset_tiled_tiff_advanced/") + copytree( + "./testdata/simple_wk_dataset/", + "./testoutput/chunking_dataset_tiled_tiff_advanced/", + ) + + ds = WKDataset("./testoutput/chunking_dataset_tiled_tiff_advanced/") + view = ds.get_view( + "color", "1", size=(150, 150, 54), offset=(10, 10, 10), is_bounded=False + ) + + for_each_chunking_advanced(ds, view) + + def test_tiled_tiff_inverse_pattern(): delete_dir("./testoutput/tiled_tiff_dataset_inverse") tiled_tiff_ds = TiledTiffDataset.create( @@ -699,7 +934,7 @@ def test_tiled_tiff_inverse_pattern(): pattern="{zzz}/{xxx}/{yyy}.tif", ) - mag = tiled_tiff_ds.add_layer("color", "color").add_mag("1") + mag = tiled_tiff_ds.add_layer("color", Layer.COLOR_TYPE).add_mag("1") data = np.zeros((250, 200, 10), dtype=np.uint8) for h in range(10): @@ -738,7 +973,7 @@ def test_view_write_without_open(): ds.get_layer("color").add_mag("1") - wk_view = ds.get_view("color", "1", size=(32, 64, 16)) + wk_view = ds.get_view("color", "1", size=(32, 64, 16), is_bounded=False) assert not wk_view._is_opened diff --git a/wkcuber/api/Dataset.py b/wkcuber/api/Dataset.py index 73c449c4e..782083a2d 100644 --- a/wkcuber/api/Dataset.py +++ b/wkcuber/api/Dataset.py @@ -112,12 +112,13 @@ def delete_layer(self, layer_name): # delete files on disk rmtree(join(self.path, layer_name)) - def get_view(self, layer_name, mag_name, size, global_offset=(0, 0, 0)) -> View: + def get_view( + self, layer_name, mag_name, size, offset=(0, 0, 0), is_bounded=True + ) -> View: layer = self.get_layer(layer_name) mag = layer.get_mag(mag_name) - mag_file_path = path.join(self.path, layer.name, mag.name) - return mag.get_view(mag_file_path, size=size, global_offset=global_offset) + return mag.get_view(size=size, offset=offset, is_bounded=is_bounded) def _create_layer(self, layer_name, dtype, num_channels) -> Layer: raise NotImplementedError diff --git a/wkcuber/api/MagDataset.py b/wkcuber/api/MagDataset.py index 8684fd130..3271a8cea 100644 --- a/wkcuber/api/MagDataset.py +++ b/wkcuber/api/MagDataset.py @@ -14,13 +14,7 @@ def __init__(self, layer, name): self.name = name self.header = self.get_header() - file_path = join(self.layer.dataset.path, self.layer.name, self.name) - size = self.layer.dataset.properties.data_layers[ - self.layer.name - ].get_bounding_box_size() - self.view = self.get_view( - file_path, size, global_offset=(0, 0, 0), is_bounded=False - ) + self.view = self.get_view(is_bounded=False) def open(self): self.view.initialize() @@ -53,7 +47,38 @@ def write(self, data, offset=(0, 0, 0)): def get_header(self): raise NotImplementedError - def get_view(self, mag_file_path, size, global_offset, is_bounded=True): + def get_view(self, size=None, offset=None, is_bounded=True): + size_in_properties = self.layer.dataset.properties.data_layers[ + self.layer.name + ].get_bounding_box_size() + + if offset is None: + offset = (0, 0, 0) + + if size is None: + size = size_in_properties + + # assert that the parameter size is valid + for s1, s2, off in zip(size_in_properties, size, offset): + if s2 + off > s1 and is_bounded: + raise AssertionError( + f"The combination of the passed parameter 'size' {size} and {offset} are not compatible with the " + f"size ({size_in_properties}) from the properties.json." + ) + + mag_file_path = join(self.layer.dataset.path, self.layer.name, self.name) + offset_in_properties = self.layer.dataset.properties.data_layers[ + self.layer.name + ].get_bounding_box_offset() + dataset_offset = ( + (0, 0, 0) if offset_in_properties == (-1, -1, -1) else offset_in_properties + ) + global_offset = np.array(dataset_offset) + np.array(offset) + return self._get_view_type()( + mag_file_path, self.header, size, global_offset, is_bounded + ) + + def _get_view_type(self): raise NotImplementedError def _assert_valid_num_channels(self, write_data_shape): @@ -85,9 +110,6 @@ def get_header(self) -> wkw.Header: block_type=self.block_type, ) - def get_view(self, mag_file_path, size, global_offset, is_bounded=True) -> WKView: - return WKView(mag_file_path, self.header, size, global_offset, is_bounded) - @classmethod def create(cls, layer, name, block_len, file_len, block_type): mag_dataset = cls(layer, name, block_len, file_len, block_type) @@ -97,6 +119,9 @@ def create(cls, layer, name, block_len, file_len, block_type): return mag_dataset + def _get_view_type(self): + return WKView + class TiffMagDataset(MagDataset): def __init__(self, layer, name, pattern): @@ -111,14 +136,14 @@ def get_header(self) -> TiffMagHeader: tile_size=self.layer.dataset.properties.tile_size, ) - def get_view(self, mag_file_path, size, global_offset, is_bounded=True) -> TiffView: - return TiffView(mag_file_path, self.header, size, global_offset, is_bounded) - @classmethod def create(cls, layer, name, pattern): mag_dataset = cls(layer, name, pattern) return mag_dataset + def _get_view_type(self): + return TiffView + class TiledTiffMagDataset(TiffMagDataset): def get_tile(self, x_index, y_index, z_index) -> np.array: diff --git a/wkcuber/api/View.py b/wkcuber/api/View.py index 0ae41b183..73b9b5b7b 100644 --- a/wkcuber/api/View.py +++ b/wkcuber/api/View.py @@ -1,7 +1,11 @@ +import math + import numpy as np from wkw import Dataset from wkcuber.api.TiffData.TiffMag import TiffMag +from wkcuber.api.bounding_box import BoundingBox +from wkcuber.utils import wait_and_ensure_success class View: @@ -69,6 +73,17 @@ def read(self, size=None, offset=(0, 0, 0)) -> np.array: return data + def get_view(self, size, offset=(0, 0, 0)): + self.assert_bounds(offset, size) + view_offset = self.global_offset + np.array(offset) + return type(self)( + self.path, + self.header, + size=size, + global_offset=view_offset, + is_bounded=self.is_bounded, + ) + def check_bounds(self, offset, size) -> bool: for s1, s2, off in zip(self.size, size, offset): if s2 + off > s1 and self.is_bounded: @@ -81,6 +96,28 @@ def assert_bounds(self, offset, size): f"Writing out of bounds: The passed parameter 'size' {size} exceeds the size of the current view ({self.size})" ) + def for_each_chunk(self, work_on_chunk, job_args_per_chunk, chunk_size, executor): + self._check_chunk_size(chunk_size) + + job_args = [] + + for chunk in BoundingBox(self.global_offset, self.size).chunk( + chunk_size, chunk_size + ): + relative_offset = np.array(chunk.topleft) - np.array(self.global_offset) + job_args.append( + ( + self.get_view(size=chunk.size, offset=relative_offset), + job_args_per_chunk, + ) + ) + + # execute the work for each chunk + wait_and_ensure_success(executor.map_to_futures(work_on_chunk, job_args)) + + def _check_chunk_size(self, chunk_size): + raise NotImplementedError + def __enter__(self): return self @@ -99,6 +136,24 @@ def open(self): self._is_opened = True return self + def _check_chunk_size(self, chunk_size): + assert chunk_size is not None + + if 0 in chunk_size: + raise AssertionError( + f"The passed parameter 'chunk_size' {chunk_size} contains at least one 0. This is not allowed." + ) + if not np.all( + np.array([math.log2(size).is_integer() for size in np.array(chunk_size)]) + ): + raise AssertionError( + f"Each element of the passed parameter 'chunk_size' {chunk_size} must be a power of 2.." + ) + if (np.array(chunk_size) % (32, 32, 32)).any(): + raise AssertionError( + f"The passed parameter 'chunk_size' {chunk_size} must be a multiple of (32, 32, 32)." + ) + class TiffView(View): def open(self): @@ -109,6 +164,30 @@ def open(self): self._is_opened = True return self + def _check_chunk_size(self, chunk_size): + assert chunk_size is not None + + if 0 in chunk_size: + raise AssertionError( + f"The passed parameter 'chunk_size' {chunk_size} contains at least one 0. This is not allowed." + ) + + if self.header.tile_size is None: + # non tiled tiff dataset + if self.size[0:2] != chunk_size[0:2]: + raise AssertionError( + f"The x- and y-length of the passed parameter 'chunk_size' {chunk_size} do not match with the size of the view {self.size}." + ) + else: + # tiled tiff dataset + file_dim = tuple(self.header.tile_size) + ( + 1, + ) # the z-dimension of an image is 1 + if (np.array(chunk_size) % file_dim).any(): + raise AssertionError( + f"The passed parameter 'chunk_size' {chunk_size} must be a multiple of the file size {file_dim}" + ) + def assert_non_negative_offset(offset): all_positive = all(i >= 0 for i in offset) diff --git a/wkcuber/api/bounding_box.py b/wkcuber/api/bounding_box.py new file mode 100644 index 000000000..3ca68a1c8 --- /dev/null +++ b/wkcuber/api/bounding_box.py @@ -0,0 +1,214 @@ +from typing import Tuple, Union, List, Dict, Generator, Optional +import json +import numpy as np +from wkcuber.mag import Mag + +Shape3D = Union[List[int], Tuple[int, int, int], np.ndarray] + + +class BoundingBox: + def __init__(self, topleft: Shape3D, size: Shape3D): + + self.topleft = np.array(topleft, dtype=np.int) + self.size = np.array(size, dtype=np.int) + + @property + def bottomright(self): + + return self.topleft + self.size + + @staticmethod + def from_wkw(bbox: Dict): + return BoundingBox( + bbox["topLeft"], [bbox["width"], bbox["height"], bbox["depth"]] + ) + + @staticmethod + def from_config(bbox: Dict): + return BoundingBox(bbox["topleft"], bbox["size"]) + + @staticmethod + def from_tuple6(tuple6: Tuple): + return BoundingBox(tuple6[0:3], tuple6[3:6]) + + @staticmethod + def from_tuple2(tuple2: Tuple): + return BoundingBox(tuple2[0], tuple2[1]) + + @staticmethod + def from_auto(obj): + if isinstance(obj, BoundingBox): + return obj + elif isinstance(obj, str): + return BoundingBox.from_auto(json.loads(obj)) + elif isinstance(obj, dict): + return BoundingBox.from_wkw(obj) + elif isinstance(obj, list) or isinstance(obj, tuple): + if len(obj) == 2: + return BoundingBox.from_tuple2(obj) + elif len(obj) == 6: + return BoundingBox.from_tuple6(obj) + + raise Exception("Unknown bounding box format.") + + def as_tuple2_string(self): + return str([self.topleft, self.size]) + + def as_wkw(self): + + width, height, depth = self.size.tolist() + + return { + "topLeft": self.topleft.tolist(), + "width": width, + "height": height, + "depth": depth, + } + + def as_config(self): + + return {"topleft": self.topleft.tolist(), "size": self.size.tolist()} + + def as_checkpoint_name(self): + + x, y, z = self.topleft + width, height, depth = self.size + return "{x}_{y}_{z}_{width}_{height}_{depth}".format( + x=x, y=y, z=z, width=width, height=height, depth=depth + ) + + def __repr__(self): + + return "BoundingBox(topleft={}, size={})".format( + str(tuple(self.topleft)), str(tuple(self.size)) + ) + + def __str__(self): + + return self.__repr__() + + def __eq__(self, other): + + return np.array_equal(self.topleft, other.topleft) and np.array_equal( + self.size, other.size + ) + + def padded_with_margins( + self, margins_left: Shape3D, margins_right: Optional[Shape3D] = None + ) -> "BoundingBox": + + if margins_right is None: + margins_right = margins_left + + margins_left = np.array(margins_left) + margins_right = np.array(margins_right) + + return BoundingBox( + topleft=self.topleft - margins_left, + size=self.size + (margins_left + margins_right), + ) + + def intersected_with( + self, other: "BoundingBox", dont_assert=False + ) -> "BoundingBox": + """ If dont_assert is set to False, this method may return empty bounding boxes (size == (0, 0, 0)) """ + + topleft = np.maximum(self.topleft, other.topleft) + bottomright = np.minimum(self.bottomright, other.bottomright) + size = np.maximum(bottomright - topleft, (0, 0, 0)) + + intersection = BoundingBox(topleft, size) + + if not dont_assert: + assert ( + not intersection.is_empty() + ), "No intersection between bounding boxes {} and {}.".format(self, other) + + return intersection + + def extended_by(self, other: "BoundingBox"): + + topleft = np.minimum(self.topleft, other.topleft) + bottomright = np.maximum(self.bottomright, other.bottomright) + size = bottomright - topleft + + return BoundingBox(topleft, size) + + def is_empty(self) -> bool: + + return not all(self.size > 0) + + def in_mag(self, mag: Mag) -> "BoundingBox": + + np_mag = np.array(mag.to_array()) + + return BoundingBox( + topleft=(self.topleft / np_mag).astype(np.int), + size=(self.size / np_mag).astype(np.int), + ) + + def contains(self, coord: Shape3D) -> bool: + + coord = np.array(coord) + + return np.all(coord >= self.topleft) and np.all( + coord < self.topleft + self.size + ) + + def chunk( + self, chunk_size: Shape3D, chunk_border_alignments: Optional[List[int]] = None + ) -> Generator["BoundingBox", None, None]: + """Decompose the bounding box into smaller chunks of size `chunk_size`. + + Chunks at the border of the bounding box might be smaller than chunk_size. + If `chunk_border_alignment` is set, all border coordinates + *between two chunks* will be divisible by that value. + """ + + start = self.topleft.copy() + chunk_size = np.array(chunk_size) + + start_adjust = np.array([0, 0, 0]) + if chunk_border_alignments is not None: + + chunk_border_alignments = np.array(chunk_border_alignments) + assert np.all( + chunk_size % chunk_border_alignments == 0 + ), "{} not divisible by {}".format(chunk_size, chunk_border_alignments) + + # Move the start to be aligned correctly. This doesn't actually change + # the start of the first chunk, because we'll intersect with `self`, + # but it'll lead to all chunk borders being aligned correctly. + start_adjust = start % chunk_border_alignments + + for x in range( + start[0] - start_adjust[0], start[0] + self.size[0], chunk_size[0] + ): + for y in range( + start[1] - start_adjust[1], start[1] + self.size[1], chunk_size[1] + ): + for z in range( + start[2] - start_adjust[2], start[2] + self.size[2], chunk_size[2] + ): + + yield BoundingBox([x, y, z], chunk_size).intersected_with(self) + + def volume(self) -> int: + + return self.size.prod() + + def slice_array(self, array: np.ndarray) -> np.ndarray: + + return array[ + self.topleft[0] : self.bottomright[0], + self.topleft[1] : self.bottomright[1], + self.topleft[2] : self.bottomright[2], + ] + + def copy(self) -> "BoundingBox": + + return BoundingBox(self.topleft.copy(), self.bottomright.copy()) + + def offset(self, vector: Tuple[int, int, int]) -> "BoundingBox": + + return BoundingBox(self.topleft + np.array(vector), self.size.copy())