From a5e65706d012c90e7e0d7908ec677699c59d8b62 Mon Sep 17 00:00:00 2001 From: Can Zhao Date: Tue, 24 May 2022 13:53:43 -0400 Subject: [PATCH 1/2] add box flip transforms Signed-off-by: Can Zhao --- monai/apps/detection/transforms/array.py | 32 +++- monai/apps/detection/transforms/box_ops.py | 42 +++- monai/apps/detection/transforms/dictionary.py | 181 ++++++++++++++++-- tests/test_box_transform.py | 38 +++- 4 files changed, 265 insertions(+), 28 deletions(-) diff --git a/monai/apps/detection/transforms/array.py b/monai/apps/detection/transforms/array.py index 901ed60615..2755635852 100644 --- a/monai/apps/detection/transforms/array.py +++ b/monai/apps/detection/transforms/array.py @@ -13,7 +13,7 @@ https://github.com/Project-MONAI/MONAI/wiki/MONAI_Design """ -from typing import Sequence, Type, Union +from typing import Optional, Sequence, Type, Union import numpy as np @@ -23,9 +23,9 @@ from monai.utils import ensure_tuple, ensure_tuple_rep, fall_back_tuple, look_up_option from monai.utils.enums import TransformBackends -from .box_ops import apply_affine_to_boxes, resize_boxes, zoom_boxes +from .box_ops import apply_affine_to_boxes, flip_boxes, resize_boxes, zoom_boxes -__all__ = ["ConvertBoxToStandardMode", "ConvertBoxMode", "AffineBox", "ZoomBox", "ResizeBox"] +__all__ = ["ConvertBoxToStandardMode", "ConvertBoxMode", "AffineBox", "ZoomBox", "ResizeBox", "FlipBox"] class ConvertBoxMode(Transform): @@ -264,3 +264,29 @@ def __call__(self, boxes: NdarrayOrTensor, src_spatial_size: Union[Sequence[int] spatial_size_ = tuple(int(round(s * scale)) for s in src_spatial_size_) return resize_boxes(boxes, src_spatial_size_, spatial_size_) + + +class FlipBox(Transform): + """ + Reverses the box coordinates along the given spatial axis. Preserves shape. + We suggest performing BoxClipToImage before this transform. + + Args: + spatial_axis: spatial axes along which to flip over. Default is None. + The default `axis=None` will flip over all of the axes of the input array. + If axis is negative it counts from the last to the first axis. + If axis is a tuple of ints, flipping is performed on all of the axes + specified in the tuple. + + """ + + backend = [TransformBackends.TORCH, TransformBackends.NUMPY] + + def __init__(self, spatial_axis: Optional[Union[Sequence[int], int]] = None) -> None: + self.spatial_axis = spatial_axis + + def __call__( # type: ignore + self, boxes: NdarrayOrTensor, spatial_size: Union[Sequence[int], int] + ) -> NdarrayOrTensor: + + return flip_boxes(boxes, spatial_size=spatial_size, flip_axes=self.spatial_axis) diff --git a/monai/apps/detection/transforms/box_ops.py b/monai/apps/detection/transforms/box_ops.py index 13911a837a..7800edcbf0 100644 --- a/monai/apps/detection/transforms/box_ops.py +++ b/monai/apps/detection/transforms/box_ops.py @@ -9,14 +9,15 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Sequence, Union +from copy import deepcopy +from typing import Optional, Sequence, Union import torch from monai.config.type_definitions import NdarrayOrTensor -from monai.data.box_utils import get_spatial_dims +from monai.data.box_utils import TO_REMOVE, get_spatial_dims from monai.transforms.utils import create_scale -from monai.utils.misc import ensure_tuple_rep +from monai.utils.misc import ensure_tuple, ensure_tuple_rep from monai.utils.type_conversion import convert_data_type, convert_to_dst_type @@ -151,3 +152,38 @@ def resize_boxes( zoom = [dst_spatial_size[axis] / float(src_spatial_size[axis]) for axis in range(spatial_dims)] return zoom_boxes(boxes=boxes, zoom=zoom) + + +def flip_boxes( + boxes: NdarrayOrTensor, + spatial_size: Union[Sequence[int], int], + flip_axes: Optional[Union[Sequence[int], int]] = None, +) -> NdarrayOrTensor: + """ + Flip boxes when the corresponding image is flipped + + Args: + boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` + spatial_size: image spatial size. + flip_axes: spatial axes along which to flip over. Default is None. + The default `axis=None` will flip over all of the axes of the input array. + If axis is negative it counts from the last to the first axis. + If axis is a tuple of ints, flipping is performed on all of the axes + specified in the tuple. + + Returns: + flipped boxes, with same data type as ``boxes``, does not share memory with ``boxes`` + """ + spatial_dims: int = get_spatial_dims(boxes=boxes) + spatial_size = ensure_tuple_rep(spatial_size, spatial_dims) + if flip_axes is None: + flip_axes = tuple(range(0, spatial_dims)) + flip_axes = ensure_tuple(flip_axes) + + # flip box + flip_boxes = deepcopy(boxes) + for axis in flip_axes: + flip_boxes[:, axis + spatial_dims] = spatial_size[axis] - boxes[:, axis] - TO_REMOVE + flip_boxes[:, axis] = spatial_size[axis] - boxes[:, axis + spatial_dims] - TO_REMOVE + + return flip_boxes diff --git a/monai/apps/detection/transforms/dictionary.py b/monai/apps/detection/transforms/dictionary.py index 41acc34149..7f57de3943 100644 --- a/monai/apps/detection/transforms/dictionary.py +++ b/monai/apps/detection/transforms/dictionary.py @@ -22,12 +22,12 @@ import numpy as np import torch -from monai.apps.detection.transforms.array import AffineBox, ConvertBoxMode, ConvertBoxToStandardMode, ZoomBox +from monai.apps.detection.transforms.array import AffineBox, ConvertBoxMode, ConvertBoxToStandardMode, FlipBox, ZoomBox from monai.config import KeysCollection from monai.config.type_definitions import NdarrayOrTensor from monai.data.box_utils import BoxMode from monai.data.utils import orientation_ras_lps -from monai.transforms import RandZoom, SpatialPad, Zoom +from monai.transforms import Flip, RandFlip, RandZoom, SpatialPad, Zoom from monai.transforms.inverse import InvertibleTransform from monai.transforms.transform import MapTransform, RandomizableTransform from monai.utils import InterpolateMode, NumpyPadMode, PytorchPadMode, ensure_tuple, ensure_tuple_rep @@ -50,6 +50,12 @@ "RandZoomBoxd", "RandZoomBoxD", "RandZoomBoxDict", + "FlipBoxd", + "FlipBoxD", + "FlipBoxDict", + "RandFlipBoxd", + "RandFlipBoxD", + "RandFlipBoxDict", ] DEFAULT_POST_FIX = PostFix.meta() @@ -249,7 +255,7 @@ def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, Nd class ZoomBoxd(MapTransform, InvertibleTransform): """ - Zooms input arrays with given probability within given zoom range. + Dictionary-based transfrom that zooms input boxes and images with the given zoom scale. Args: image_keys: Keys to pick image data for transformation. @@ -295,13 +301,12 @@ def __init__( self.image_keys = ensure_tuple(image_keys) self.box_keys = ensure_tuple(box_keys) super().__init__(self.image_keys + self.box_keys, allow_missing_keys) + self.box_ref_image_keys = ensure_tuple_rep(box_ref_image_keys, len(self.box_keys)) self.mode = ensure_tuple_rep(mode, len(self.image_keys)) self.padding_mode = ensure_tuple_rep(padding_mode, len(self.image_keys)) self.align_corners = ensure_tuple_rep(align_corners, len(self.image_keys)) self.zoomer = Zoom(zoom=zoom, keep_size=keep_size, **kwargs) - - self.box_ref_image_keys = ensure_tuple_rep(box_ref_image_keys, len(self.box_keys)) self.keep_size = keep_size def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: @@ -363,8 +368,6 @@ def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, Nd ) # Size might be out by 1 voxel so pad d[key] = SpatialPad(transform[TraceKeys.EXTRA_INFO]["original_shape"], mode="edge")(d[key]) - # Remove the applied transform - self.pop_transform(d, key) # zoom boxes if key_type == "box_key": @@ -372,15 +375,16 @@ def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, Nd src_spatial_size = transform[TraceKeys.EXTRA_INFO]["src_spatial_size"] box_inverse_transform = ZoomBox(zoom=(1 / zoom).tolist(), keep_size=self.zoomer.keep_size) d[key] = box_inverse_transform(d[key], src_spatial_size=src_spatial_size) - # Remove the applied transform - self.pop_transform(d, key) + + # Remove the applied transform + self.pop_transform(d, key) return d class RandZoomBoxd(RandomizableTransform, MapTransform, InvertibleTransform): """ - Randomly zooms input arrays with given probability within given zoom range. + Dictionary-based transfrom that randomly zooms input boxes and images with given probability within given zoom range. Args: image_keys: Keys to pick image data for transformation. @@ -439,13 +443,12 @@ def __init__( self.box_keys = ensure_tuple(box_keys) MapTransform.__init__(self, self.image_keys + self.box_keys, allow_missing_keys) RandomizableTransform.__init__(self, prob) + self.box_ref_image_keys = ensure_tuple_rep(box_ref_image_keys, len(self.box_keys)) self.rand_zoom = RandZoom(prob=1.0, min_zoom=min_zoom, max_zoom=max_zoom, keep_size=keep_size, **kwargs) self.mode = ensure_tuple_rep(mode, len(self.image_keys)) self.padding_mode = ensure_tuple_rep(padding_mode, len(self.image_keys)) self.align_corners = ensure_tuple_rep(align_corners, len(self.image_keys)) - - self.box_ref_image_keys = ensure_tuple_rep(box_ref_image_keys, len(self.box_keys)) self.keep_size = keep_size def set_random_state( @@ -508,14 +511,11 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: d = deepcopy(dict(data)) - # zoom image, copied from monai.transforms.spatial.dictionary.RandZoomd for key in self.key_iterator(d): transform = self.get_most_recent_transform(d, key) key_type = transform[TraceKeys.EXTRA_INFO]["type"] # Check if random transform was actually performed (based on `prob`) if transform[TraceKeys.DO_TRANSFORM]: - # Create inverse transform - # zoom image, copied from monai.transforms.spatial.dictionary.Zoomd if key_type == "image_key": zoom = np.array(transform[TraceKeys.EXTRA_INFO]["zoom"]) @@ -531,8 +531,6 @@ def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, Nd ) # Size might be out by 1 voxel so pad d[key] = SpatialPad(transform[TraceKeys.EXTRA_INFO]["original_shape"], mode="edge")(d[key]) - # Remove the applied transform - self.pop_transform(d, key) # zoom boxes if key_type == "box_key": @@ -541,8 +539,151 @@ def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, Nd src_spatial_size = transform[TraceKeys.EXTRA_INFO]["src_spatial_size"] box_inverse_transform = ZoomBox(zoom=(1.0 / zoom).tolist(), keep_size=self.rand_zoom.keep_size) d[key] = box_inverse_transform(d[key], src_spatial_size=src_spatial_size) - # Remove the applied transform - self.pop_transform(d, key) + + # Remove the applied transform + self.pop_transform(d, key) + return d + + +class FlipBoxd(MapTransform, InvertibleTransform): + """ + Dictionary-based transfrom that flip boxes and images. + + Args: + image_keys: Keys to pick image data for transformation. + box_keys: Keys to pick box data for transformation. The box mode is assumed to be ``StandardMode``. + box_ref_image_keys: Keys that represents the reference images to which ``box_keys`` are attached. + spatial_axis: Spatial axes along which to flip over. Default is None. + allow_missing_keys: don't raise exception if key is missing. + """ + + backend = Flip.backend + + def __init__( + self, + image_keys: KeysCollection, + box_keys: KeysCollection, + box_ref_image_keys: KeysCollection, + spatial_axis: Optional[Union[Sequence[int], int]] = None, + allow_missing_keys: bool = False, + ) -> None: + self.image_keys = ensure_tuple(image_keys) + self.box_keys = ensure_tuple(box_keys) + super().__init__(self.image_keys + self.box_keys, allow_missing_keys) + self.box_ref_image_keys = ensure_tuple_rep(box_ref_image_keys, len(self.box_keys)) + + self.flipper = Flip(spatial_axis=spatial_axis) + self.box_flipper = FlipBox(spatial_axis=self.flipper.spatial_axis) + + def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + d = dict(data) + + for key in self.image_keys: + d[key] = self.flipper(d[key]) + self.push_transform(d, key, extra_info={"type": "image_key"}) + + for box_key, box_ref_image_key in zip(self.box_keys, self.box_ref_image_keys): + spatial_size = d[box_ref_image_key].shape[1:] + d[box_key] = self.box_flipper(d[box_key], spatial_size) + self.push_transform(d, box_key, extra_info={"spatial_size": spatial_size, "type": "box_key"}) + return d + + def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + d = deepcopy(dict(data)) + + for key in self.key_iterator(d): + transform = self.get_most_recent_transform(d, key) + key_type = transform[TraceKeys.EXTRA_INFO]["type"] + + # flip image, copied from monai.transforms.spatial.dictionary.Flipd + if key_type == "image_key": + d[key] = self.flipper(d[key]) + + # flip boxes + if key_type == "box_key": + spatial_size = transform[TraceKeys.EXTRA_INFO]["spatial_size"] + d[key] = self.box_flipper(d[key], spatial_size) + + # Remove the applied transform + self.pop_transform(d, key) + return d + + +class RandFlipBoxd(RandomizableTransform, MapTransform, InvertibleTransform): + """ + Dictionary-based transfrom that randomly flip boxes and images with the given probabilities. + + Args: + image_keys: Keys to pick image data for transformation. + box_keys: Keys to pick box data for transformation. The box mode is assumed to be ``StandardMode``. + box_ref_image_keys: Keys that represents the reference images to which ``box_keys`` are attached. + prob: Probability of flipping. + spatial_axis: Spatial axes along which to flip over. Default is None. + allow_missing_keys: don't raise exception if key is missing. + """ + + backend = RandFlip.backend + + def __init__( + self, + image_keys: KeysCollection, + box_keys: KeysCollection, + box_ref_image_keys: KeysCollection, + prob: float = 0.1, + spatial_axis: Optional[Union[Sequence[int], int]] = None, + allow_missing_keys: bool = False, + ) -> None: + self.image_keys = ensure_tuple(image_keys) + self.box_keys = ensure_tuple(box_keys) + MapTransform.__init__(self, self.image_keys + self.box_keys, allow_missing_keys) + RandomizableTransform.__init__(self, prob) + self.box_ref_image_keys = ensure_tuple_rep(box_ref_image_keys, len(self.box_keys)) + + self.flipper = RandFlip(prob=1.0, spatial_axis=spatial_axis) + self.box_flipper = FlipBox(spatial_axis=spatial_axis) + + def set_random_state( + self, seed: Optional[int] = None, state: Optional[np.random.RandomState] = None + ) -> "RandFlipBoxd": + super().set_random_state(seed, state) + self.flipper.set_random_state(seed, state) + return self + + def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + d = dict(data) + self.randomize(None) + + for key in self.image_keys: + if self._do_transform: + d[key] = self.flipper(d[key], randomize=False) + self.push_transform(d, key, extra_info={"type": "image_key"}) + + for box_key, box_ref_image_key in zip(self.box_keys, self.box_ref_image_keys): + spatial_size = d[box_ref_image_key].shape[1:] + if self._do_transform: + d[box_key] = self.box_flipper(d[box_key], spatial_size) + self.push_transform(d, box_key, extra_info={"spatial_size": spatial_size, "type": "box_key"}) + return d + + def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]: + d = deepcopy(dict(data)) + + for key in self.key_iterator(d): + transform = self.get_most_recent_transform(d, key) + key_type = transform[TraceKeys.EXTRA_INFO]["type"] + # Check if random transform was actually performed (based on `prob`) + if transform[TraceKeys.DO_TRANSFORM]: + # flip image, copied from monai.transforms.spatial.dictionary.RandFlipd + if key_type == "image_key": + d[key] = self.flipper(d[key], randomize=False) + + # flip boxes + if key_type == "box_key": + spatial_size = transform[TraceKeys.EXTRA_INFO]["spatial_size"] + d[key] = self.box_flipper(d[key], spatial_size) + + # Remove the applied transform + self.pop_transform(d, key) return d @@ -551,3 +692,5 @@ def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, Nd ZoomBoxD = ZoomBoxDict = ZoomBoxd RandZoomBoxD = RandZoomBoxDict = RandZoomBoxd AffineBoxToImageCoordinateD = AffineBoxToImageCoordinateDict = AffineBoxToImageCoordinated +FlipBoxD = FlipBoxDict = FlipBoxd +RandFlipBoxD = RandFlipBoxDict = RandFlipBoxd diff --git a/tests/test_box_transform.py b/tests/test_box_transform.py index 10e311971e..447dd2e749 100644 --- a/tests/test_box_transform.py +++ b/tests/test_box_transform.py @@ -18,6 +18,8 @@ from monai.apps.detection.transforms.dictionary import ( AffineBoxToImageCoordinated, ConvertBoxModed, + FlipBoxd, + RandFlipBoxd, RandZoomBoxd, ZoomBoxd, ) @@ -38,8 +40,8 @@ p([[0, 0, 0, 0, 0, 0], [0, 1, 0, 2, 2, 3], [0, 1, 1, 2, 2, 3]]), p([[0, 0, 0, 0, 0, 0], [0, 3, 0, 1, 9, 4.5], [0, 3, 1.5, 1, 9, 6]]), p([[1, -6, -1, 1, -6, -1], [1, -3, -1, 2, 3, 3.5], [1, -3, 0.5, 2, 3, 5]]), + p([[4, 6, 4, 4, 6, 4], [2, 3, 1, 4, 5, 4], [2, 3, 0, 4, 5, 3]]), p([[0, 1, 0, 2, 3, 3], [0, 1, 1, 2, 3, 4]]), - p([[0, 1, 1, 2, 3, 4], [0, 1, 0, 2, 3, 3]]), ] ) @@ -53,8 +55,8 @@ def test_value( expected_convert_result, expected_zoom_result, expected_zoom_keepsize_result, - expected_clip_result, expected_flip_result, + expected_clip_result, ): test_dtype = [torch.float32] for dtype in test_dtype: @@ -93,7 +95,7 @@ def test_value( zoom_result["boxes"], expected_zoom_keepsize_result, type_test=True, device_test=True, atol=1e-3 ) - # test ZoomBoxd + # test RandZoomBoxd transform_zoom = RandZoomBoxd( image_keys="image", box_keys="boxes", @@ -124,6 +126,36 @@ def test_value( data_back = invert_transform_affine(affine_result) assert_allclose(data_back["boxes"], data["boxes"], type_test=False, device_test=False, atol=0.01) + # test FlipBoxd + transform_flip = FlipBoxd( + image_keys="image", box_keys="boxes", box_ref_image_keys="image", spatial_axis=[0, 1, 2] + ) + flip_result = transform_flip(data) + assert_allclose(flip_result["boxes"], expected_flip_result, type_test=True, device_test=True, atol=1e-3) + invert_transform_flip = Invertd( + keys=["image", "boxes"], transform=transform_flip, orig_keys=["image", "boxes"] + ) + data_back = invert_transform_flip(flip_result) + assert_allclose(data_back["boxes"], data["boxes"], type_test=False, device_test=False, atol=1e-3) + assert_allclose(data_back["image"], data["image"], type_test=False, device_test=False, atol=1e-3) + + # test RandFlipBoxd + for spatial_axis in [(0,), (1,), (2,), (0, 1), (1, 2)]: + transform_flip = RandFlipBoxd( + image_keys="image", + box_keys="boxes", + box_ref_image_keys="image", + prob=1.0, + spatial_axis=spatial_axis, + ) + flip_result = transform_flip(data) + invert_transform_flip = Invertd( + keys=["image", "boxes"], transform=transform_flip, orig_keys=["image", "boxes"] + ) + data_back = invert_transform_flip(flip_result) + assert_allclose(data_back["boxes"], data["boxes"], type_test=False, device_test=False, atol=1e-3) + assert_allclose(data_back["image"], data["image"], type_test=False, device_test=False, atol=1e-3) + if __name__ == "__main__": unittest.main() From 92248040a2e16243a571b5d8a6cc4874dfe9bbd5 Mon Sep 17 00:00:00 2001 From: Can Zhao Date: Tue, 24 May 2022 14:00:15 -0400 Subject: [PATCH 2/2] update docstring Signed-off-by: Can Zhao --- monai/apps/detection/transforms/array.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/monai/apps/detection/transforms/array.py b/monai/apps/detection/transforms/array.py index 2755635852..131385a142 100644 --- a/monai/apps/detection/transforms/array.py +++ b/monai/apps/detection/transforms/array.py @@ -269,7 +269,6 @@ def __call__(self, boxes: NdarrayOrTensor, src_spatial_size: Union[Sequence[int] class FlipBox(Transform): """ Reverses the box coordinates along the given spatial axis. Preserves shape. - We suggest performing BoxClipToImage before this transform. Args: spatial_axis: spatial axes along which to flip over. Default is None. @@ -288,5 +287,10 @@ def __init__(self, spatial_axis: Optional[Union[Sequence[int], int]] = None) -> def __call__( # type: ignore self, boxes: NdarrayOrTensor, spatial_size: Union[Sequence[int], int] ) -> NdarrayOrTensor: + """ + Args: + boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode`` + spatial_size: image spatial size. + """ return flip_boxes(boxes, spatial_size=spatial_size, flip_axes=self.spatial_axis)