From e2a170583b2ba24024539e711993f2797a76b937 Mon Sep 17 00:00:00 2001 From: Tim Blakely Date: Tue, 9 Aug 2022 12:02:10 -0700 Subject: [PATCH] Allow `expected_output_box` to work with N-d bounding boxes. PiperOrigin-RevId: 466427583 --- connectomics/common/bounding_box.py | 44 +++++++++++++++++++--- connectomics/common/bounding_box_test.py | 35 +++++++++++++++++ connectomics/volume/base.py | 5 +++ connectomics/volume/base_test.py | 20 ++++++++++ connectomics/volume/subvolume_processor.py | 5 ++- 5 files changed, 102 insertions(+), 7 deletions(-) diff --git a/connectomics/common/bounding_box.py b/connectomics/common/bounding_box.py index 4557e8d..f8ff3f9 100644 --- a/connectomics/common/bounding_box.py +++ b/connectomics/common/bounding_box.py @@ -357,20 +357,24 @@ def relative(self: S, def to_slice_tuple(self, start_dim: Optional[int] = None, - end_dim: Optional[int] = None) -> Tuple[slice, ...]: - """Returns slice in C-order (reverse dimension akin to ZYX). + end_dim: Optional[int] = None, + order: str = 'c') -> Tuple[slice, ...]: + """Returns slice in C or Fortran-order (XYZ). Args: - start_dim: Optional beginning dimension to begin slice (forward format - e.g. XYZ). - end_dim: Optional end dimension to begin slice (forward format). + start_dim: Optional beginning dimension to begin slice. + end_dim: Optional end dimension to begin slice. + order: C or Fortran order. Returns: Tuple corresponding to a slice expression akin to np.index_exp. """ start = self.start[start_dim:end_dim] end = self.end[start_dim:end_dim] - return tuple(slice(s, e, None) for s, e in zip(start[::-1], end[::-1])) + extents = tuple(zip(start, end)) + if order[0].lower() == 'c': + extents = extents[::-1] + return tuple(slice(s, e, None) for s, e in extents) def to_slice3d(self) -> Tuple[slice, slice, slice]: """Convenience function for 3d use. @@ -511,3 +515,31 @@ def from_json(as_json: str) -> BoundingBoxBase: [utils.is_floatlike(v) for v in bbox.size]): return FloatBoundingBox(bbox.start, bbox.size) return BoundingBox(bbox.start, bbox.size) + + +def from_slices(slices: Sequence[slice], + order='c', + inclusive=True) -> BoundingBox: + """Creates a new BoundingBox from slice extents. + + Args: + slices: Extents for the limits of the bounding box. + order: C or Fortran order of slices (CZYX... vs XYZC...). + inclusive: If the `slice`.start is considered inclusive or not. If not, box + start/end extents will be slice.start:slice.stop-1. + + Returns: + New BoundingBox with the given extents. + + Raises: + ValueError: If any of the extents is not numeric. + """ + extents = [(s.start, s.stop) for s in slices] + if np.any(None in np.array(extents)): + raise ValueError('All slices must be finite.') + if not inclusive: + extents = tuple((a, z - 1) for a, z in extents) + if order[0].lower() == 'c': + extents = extents[::-1] + start, stop = zip(*extents) + return BoundingBox(start, end=stop) diff --git a/connectomics/common/bounding_box_test.py b/connectomics/common/bounding_box_test.py index a1560e5..b224ab0 100644 --- a/connectomics/common/bounding_box_test.py +++ b/connectomics/common/bounding_box_test.py @@ -238,6 +238,7 @@ def test_to_slice(self): expected = np.index_exp[start[4]:end[4], start[3]:end[3], start[2]:end[2], start[1]:end[1], start[0]:end[0]] self.assertEqual(expected, box.to_slice_tuple()) + self.assertEqual(expected[::-1], box.to_slice_tuple(order='f')) self.assertEqual(expected[1:], box.to_slice4d()) self.assertEqual(expected[2:], box.to_slice3d()) @@ -401,6 +402,40 @@ def test_dataclass(self): new_box = Base.from_json(box.to_json()) self.assertEqual(box, new_box) + def test_from_slices(self): + slices = [ + slice(3, 7), + slice(11, 17), + slice(-3, 9), + slice(5, 55), + slice(7, 8), + ] + + # C order, inclusive + self.assertEqual( + Box(start=[3, 11, -3, 5, 7][::-1], end=[7, 17, 9, 55, 8][::-1]), + bounding_box.from_slices(slices)) + + # C order, exclusive + self.assertEqual( + Box(start=[3, 11, -3, 5, 7][::-1], end=[6, 16, 8, 54, 7][::-1]), + bounding_box.from_slices(slices, inclusive=False)) + + # Fortran order, inclusive + self.assertEqual( + Box(start=[3, 11, -3, 5, 7], end=[7, 17, 9, 55, 8]), + bounding_box.from_slices(slices, order='f')) + + # Fortran order, exclusive + self.assertEqual( + Box(start=[3, 11, -3, 5, 7], end=[6, 16, 8, 54, 7]), + bounding_box.from_slices(slices, order='f', inclusive=False)) + + # Invalid extents + with self.assertRaises(ValueError): + slices.append(slice(3, None)) + bounding_box.from_slices(slices) + if __name__ == '__main__': absltest.main() diff --git a/connectomics/volume/base.py b/connectomics/volume/base.py index 7d0977d..8f739de 100644 --- a/connectomics/volume/base.py +++ b/connectomics/volume/base.py @@ -140,5 +140,10 @@ def chunk_size(self) -> array.Tuple4i: """Backing chunk size in voxels, CZYX.""" raise NotImplementedError + def clip_box_to_volume( + self, box: bounding_box.BoundingBox) -> bounding_box.BoundingBox: + return bounding_box.BoundingBox( + box.start, end=np.minimum(box.end, self.volume_size)) + # TODO(timblakely): determine what other attributes we want to make mandatory for # all implementations and add them here. diff --git a/connectomics/volume/base_test.py b/connectomics/volume/base_test.py index ef81f30..e6edbd2 100644 --- a/connectomics/volume/base_test.py +++ b/connectomics/volume/base_test.py @@ -16,6 +16,7 @@ from absl.testing import absltest from connectomics.common import array +from connectomics.common import bounding_box from connectomics.volume import base import numpy as np import numpy.testing as npt @@ -140,6 +141,25 @@ def write_slices(self, slices, data): v[0, 1:3, 5:, :] = expected_data self.assertTrue(v.called) + def test_clip_box(self): + + class GetPointsVolume(ShimVolume): + + @property + def volume_size(self): + return [4, 5, 6] + + v = GetPointsVolume() + box = bounding_box.BoundingBox([1, 2, 3], [100, 200, 300]) + clipped = v.clip_box_to_volume(box) + self.assertNotEqual(clipped, box) + self.assertEqual(tuple(clipped.end), (4, 5, 6)) + + v = GetPointsVolume() + box = bounding_box.BoundingBox([0, 0, 0], [2, 1, 3]) + clipped = v.clip_box_to_volume(box) + self.assertEqual(clipped, box) + if __name__ == '__main__': absltest.main() diff --git a/connectomics/volume/subvolume_processor.py b/connectomics/volume/subvolume_processor.py index f0736e6..c75b2a7 100644 --- a/connectomics/volume/subvolume_processor.py +++ b/connectomics/volume/subvolume_processor.py @@ -22,6 +22,7 @@ from connectomics.common import array from connectomics.common import bounding_box +from connectomics.common import file from connectomics.volume import descriptor from connectomics.volume import subvolume import dataclasses_json @@ -53,7 +54,9 @@ class ProcessVolumeConfig(dataclasses_json.DataClassJsonMixin): """User-supplied configuration.""" # Input volume to process. - input_volume: descriptor.VolumeDescriptor + input_volume: descriptor.VolumeDescriptor = dataclasses.field( + metadata=dataclasses_json.config( + decoder=file.dataclass_loader(descriptor.VolumeDescriptor))) # Output directory to write the volumetric data, inserted automatically into # the output_volume's TensorStore spec.