diff --git a/monai/apps/detection/transforms/dictionary.py b/monai/apps/detection/transforms/dictionary.py index fc65954358..2c29937aa4 100644 --- a/monai/apps/detection/transforms/dictionary.py +++ b/monai/apps/detection/transforms/dictionary.py @@ -43,7 +43,7 @@ from monai.transforms.transform import MapTransform, Randomizable, RandomizableTransform from monai.transforms.utils import generate_pos_neg_label_crop_centers, map_binary_to_indices from monai.utils import ImageMetaKey as Key -from monai.utils import InterpolateMode, NumpyPadMode, ensure_tuple, ensure_tuple_rep +from monai.utils import InterpolateMode, NumpyPadMode, ensure_tuple, ensure_tuple_rep, fall_back_tuple from monai.utils.enums import PostFix, TraceKeys from monai.utils.type_conversion import convert_data_type @@ -1174,9 +1174,8 @@ def randomize( # type: ignore def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> List[Dict[Hashable, torch.Tensor]]: d = dict(data) - spatial_dims = len(d[self.image_keys[0]].shape) - 1 image_size = d[self.image_keys[0]].shape[1:] - self.spatial_size = ensure_tuple_rep(self.spatial_size_, spatial_dims) + self.spatial_size = fall_back_tuple(self.spatial_size_, image_size) # randomly sample crop centers boxes = d[self.box_keys] diff --git a/tests/test_box_transform.py b/tests/test_box_transform.py index 6b0a4a2b19..ad82b080f1 100644 --- a/tests/test_box_transform.py +++ b/tests/test_box_transform.py @@ -16,7 +16,12 @@ from parameterized import parameterized from monai.apps.detection.transforms.box_ops import convert_mask_to_box -from monai.apps.detection.transforms.dictionary import BoxToMaskd, ConvertBoxModed, MaskToBoxd +from monai.apps.detection.transforms.dictionary import ( + BoxToMaskd, + ConvertBoxModed, + MaskToBoxd, + RandCropBoxByPosNegLabeld, +) from monai.transforms import CastToTyped from tests.utils import TEST_NDARRAYS, assert_allclose @@ -287,6 +292,24 @@ def test_value_3d( # 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) + def test_crop_shape(self): + tt = RandCropBoxByPosNegLabeld( + image_keys=["image"], + box_keys="box", + label_keys="label", + spatial_size=[10, 7, -1], + whole_box=True, + num_samples=1, + pos=1, + neg=0, + ) + iii = { + "image": torch.rand(1, 10, 8, 7), + "box": torch.tensor(((1.0, 2.0, 3.0, 4.0, 5.0, 6.0),)), + "label": torch.tensor((1,)).long(), + } + self.assertEqual(tt(iii)[0]["image"].shape, (1, 10, 7, 7)) + if __name__ == "__main__": unittest.main()