Safely Rotate Images Without Cropping

README.md

@@ -187,6 +187,7 @@ Spatial-level transforms will simultaneously change both an input image as well
 | [RandomSizedCrop](https://albumentations.ai/docs/api_reference/augmentations/crops/transforms/#albumentations.augmentations.crops.transforms.RandomSizedCrop)                   | ✓     | ✓     | ✓      | ✓         |
 | [Resize](https://albumentations.ai/docs/api_reference/augmentations/geometric/resize/#albumentations.augmentations.geometric.resize.Resize)                                     | ✓     | ✓     | ✓      | ✓         |
 | [Rotate](https://albumentations.ai/docs/api_reference/augmentations/geometric/rotate/#albumentations.augmentations.geometric.rotate.Rotate)                                     | ✓     | ✓     | ✓      | ✓         |
+| [SafeRotate](https://albumentations.ai/docs/api_reference/augmentations/geometric/rotate/#albumentations.augmentations.geometric.rotate.SafeRotate)                             | ✓     | ✓     | ✓      | ✓         |
 | [ShiftScaleRotate](https://albumentations.ai/docs/api_reference/augmentations/geometric/transforms/#albumentations.augmentations.geometric.transforms.ShiftScaleRotate)         | ✓     | ✓     | ✓      | ✓         |
 | [SmallestMaxSize](https://albumentations.ai/docs/api_reference/augmentations/geometric/resize/#albumentations.augmentations.geometric.resize.SmallestMaxSize)                   | ✓     | ✓     | ✓      | ✓         |
 | [Transpose](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.Transpose)                                           | ✓     | ✓     | ✓      | ✓         |

albumentations/augmentations/geometric/functional.py

@@ -410,3 +410,114 @@ def perspective_keypoint(
         return keypoint_scale((x, y, angle, scale), scale_x, scale_y)
 
     return x, y, angle, scale
+
+
+@preserve_channel_dim
+def safe_rotate(
+    img: np.ndarray,
+    angle: int = 0,
+    interpolation: int = cv2.INTER_LINEAR,
+    value: int = None,
+    border_mode: int = cv2.BORDER_REFLECT_101,
+):
+
+    old_rows, old_cols = img.shape[:2]
+
+    # getRotationMatrix2D needs coordinates in reverse order (width, height) compared to shape
+    image_center = (old_cols / 2, old_rows / 2)
+
+    # Rows and columns of the rotated image (not cropped)
+    new_rows, new_cols = safe_rotate_enlarged_img_size(angle=angle, rows=old_rows, cols=old_cols)
+
+    # Rotation Matrix
+    rotation_mat = cv2.getRotationMatrix2D(image_center, angle, 1.0)
+
+    # Shift the image to create padding
+    rotation_mat[0, 2] += new_cols / 2 - image_center[0]
+    rotation_mat[1, 2] += new_rows / 2 - image_center[1]
+
+    # CV2 Transformation function
+    warp_affine_fn = _maybe_process_in_chunks(
+        cv2.warpAffine,
+        M=rotation_mat,
+        dsize=(new_cols, new_rows),
+        flags=interpolation,
+        borderMode=border_mode,
+        borderValue=value,
+    )
+
+    # rotate image with the new bounds
+    rotated_img = warp_affine_fn(img)
+
+    # Resize image back to the original size
+    resized_img = resize(img=rotated_img, height=old_rows, width=old_cols, interpolation=interpolation)
+
+    return resized_img
+
+
+def bbox_safe_rotate(bbox, angle, rows, cols):
+    old_rows = rows
+    old_cols = cols
+
+    # Rows and columns of the rotated image (not cropped)
+    new_rows, new_cols = safe_rotate_enlarged_img_size(angle=angle, rows=old_rows, cols=old_cols)
+
+    col_diff = int(np.ceil(abs(new_cols - old_cols) / 2))
+    row_diff = int(np.ceil(abs(new_rows - old_rows) / 2))
+
+    # Normalize shifts
+    norm_col_shift = col_diff / new_cols
+    norm_row_shift = row_diff / new_rows
+
+    # shift bbox
+    shifted_bbox = (
+        bbox[0] + norm_col_shift,
+        bbox[1] + norm_row_shift,
+        bbox[2] + norm_col_shift,
+        bbox[3] + norm_row_shift,
+    )
+
+    rotated_bbox = bbox_rotate(bbox=shifted_bbox, angle=angle, rows=new_rows, cols=new_cols)
+
+    # Bounding boxes are scale invariant, so this does not need to be rescaled to the old size
+    return rotated_bbox
+
+
+def keypoint_safe_rotate(keypoint, angle, rows, cols):
+    old_rows = rows
+    old_cols = cols
+
+    # Rows and columns of the rotated image (not cropped)
+    new_rows, new_cols = safe_rotate_enlarged_img_size(angle=angle, rows=old_rows, cols=old_cols)
+
+    col_diff = int(np.ceil(abs(new_cols - old_cols) / 2))
+    row_diff = int(np.ceil(abs(new_rows - old_rows) / 2))
+
+    # Shift keypoint
+    shifted_keypoint = (keypoint[0] + col_diff, keypoint[1] + row_diff, keypoint[2], keypoint[3])
+
+    # Rotate keypoint
+    rotated_keypoint = keypoint_rotate(shifted_keypoint, angle, rows=new_rows, cols=new_cols)
+
+    # Scale the keypoint
+    return keypoint_scale(rotated_keypoint, old_cols / new_cols, old_rows / new_rows)
+
+
+def safe_rotate_enlarged_img_size(angle: float, rows: int, cols: int):
+
+    deg_angle = abs(angle)
+
+    # The rotation angle
+    angle = np.deg2rad(deg_angle % 90)
+
+    # The width of the frame to contain the rotated image
+    r_cols = cols * np.cos(angle) + rows * np.sin(angle)
+
+    # The height of the frame to contain the rotated image
+    r_rows = cols * np.sin(angle) + rows * np.cos(angle)
+
+    # The above calculations work as is for 0<90 degrees, and for 90<180 the cols and rows are flipped
+    if deg_angle > 90:
+        return int(r_cols), int(r_rows)
+    else:
+        return int(r_rows), int(r_cols)

albumentations/augmentations/geometric/rotate.py

@@ -5,7 +5,7 @@
 from . import functional as F
 from ...core.transforms_interface import DualTransform, to_tuple
 
-__all__ = ["Rotate", "RandomRotate90"]
+__all__ = ["Rotate", "RandomRotate90", "SafeRotate"]
 
 
 class RandomRotate90(DualTransform):
@@ -101,3 +101,72 @@ def apply_to_keypoint(self, keypoint, angle=0, **params):
 
     def get_transform_init_args_names(self):
         return ("limit", "interpolation", "border_mode", "value", "mask_value")
+
+
+class SafeRotate(DualTransform):
+    """Rotate the input inside the input's frame by an angle selected randomly from the uniform distribution.
+
+    The resulting image may have artifacts in it. After rotation, the image may have a different aspect ratio, and
+    after resizing, it returns to its original shape with the original aspect ratio of the image. For these reason we
+    may see some artifacts.
+
+    Args:
+        limit ((int, int) or int): range from which a random angle is picked. If limit is a single int
+            an angle is picked from (-limit, limit). Default: (-90, 90)
+        interpolation (OpenCV flag): flag that is used to specify the interpolation algorithm. Should be one of:
+            cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA, cv2.INTER_LANCZOS4.
+            Default: cv2.INTER_LINEAR.
+        border_mode (OpenCV flag): flag that is used to specify the pixel extrapolation method. Should be one of:
+            cv2.BORDER_CONSTANT, cv2.BORDER_REPLICATE, cv2.BORDER_REFLECT, cv2.BORDER_WRAP, cv2.BORDER_REFLECT_101.
+            Default: cv2.BORDER_REFLECT_101
+        value (int, float, list of ints, list of float): padding value if border_mode is cv2.BORDER_CONSTANT.
+        mask_value (int, float,
+                    list of ints,
+                    list of float): padding value if border_mode is cv2.BORDER_CONSTANT applied for masks.
+        p (float): probability of applying the transform. Default: 0.5.
+
+    Targets:
+        image, mask, bboxes, keypoints
+
+    Image types:
+        uint8, float32
+    """
+
+    def __init__(
+        self,
+        limit=90,
+        interpolation=cv2.INTER_LINEAR,
+        border_mode=cv2.BORDER_REFLECT_101,
+        value=None,
+        mask_value=None,
+        always_apply=False,
+        p=0.5,
+    ):
+        super(SafeRotate, self).__init__(always_apply, p)
+        self.limit = to_tuple(limit)
+        self.interpolation = interpolation
+        self.border_mode = border_mode
+        self.value = value
+        self.mask_value = mask_value
+
+    def apply(self, img, angle=0, interpolation=cv2.INTER_LINEAR, **params):
+        return F.safe_rotate(
+            img=img, value=self.value, angle=angle, interpolation=interpolation, border_mode=self.border_mode
+        )
+
+    def apply_to_mask(self, img, angle=0, **params):
+        return F.safe_rotate(
+            img=img, value=self.mask_value, angle=angle, interpolation=cv2.INTER_NEAREST, border_mode=self.border_mode
+        )
+
+    def get_params(self):
+        return {"angle": random.uniform(self.limit[0], self.limit[1])}
+
+    def apply_to_bbox(self, bbox, angle=0, **params):
+        return F.bbox_safe_rotate(bbox=bbox, angle=angle, rows=params["rows"], cols=params["cols"])
+
+    def apply_to_keypoint(self, keypoint, angle=0, **params):
+        return F.keypoint_safe_rotate(keypoint, angle=angle, rows=params["rows"], cols=params["cols"])
+
+    def get_transform_init_args_names(self):
+        return ("limit", "interpolation", "border_mode", "value", "mask_value")

tests/test_augmentations.py

@@ -14,6 +14,7 @@
     Transpose,
     RandomRotate90,
     Rotate,
+    SafeRotate,
     ShiftScaleRotate,
     CenterCrop,
     OpticalDistortion,
@@ -190,6 +191,7 @@ def test_image_only_augmentations_with_float_values(augmentation_cls, params, fl
         [Transpose, {}],
         [RandomRotate90, {}],
         [Rotate, {}],
+        [SafeRotate, {}],
         [CoarseDropout, {"fill_value": 0, "mask_fill_value": 0}],
         [ShiftScaleRotate, {}],
         [OpticalDistortion, {}],
@@ -224,6 +226,7 @@ def test_dual_augmentations(augmentation_cls, params, image, mask):
         [Transpose, {}],
         [RandomRotate90, {}],
         [Rotate, {}],
+        [SafeRotate, {}],
         [ShiftScaleRotate, {}],
         [OpticalDistortion, {}],
         [GridDistortion, {}],
@@ -291,6 +294,7 @@ def test_imgaug_dual_augmentations(augmentation_cls, image, mask):
         [Transpose, {}],
         [RandomRotate90, {}],
         [Rotate, {}],
+        [SafeRotate, {}],
         [ShiftScaleRotate, {}],
         [OpticalDistortion, {}],
         [GridDistortion, {}],
@@ -368,6 +372,7 @@ def test_augmentations_wont_change_input(augmentation_cls, params, image, mask):
         [Transpose, {}],
         [RandomRotate90, {}],
         [Rotate, {}],
+        [SafeRotate, {}],
         [ShiftScaleRotate, {}],
         [OpticalDistortion, {}],
         [GridDistortion, {}],
@@ -433,6 +438,7 @@ def test_augmentations_wont_change_float_input(augmentation_cls, params, float_i
         [Transpose, {}],
         [RandomRotate90, {}],
         [Rotate, {}],
+        [SafeRotate, {}],
         [OpticalDistortion, {}],
         [GridDistortion, {}],
         [ElasticTransform, {}],
@@ -504,6 +510,7 @@ def test_augmentations_wont_change_shape_grayscale(augmentation_cls, params, ima
         [Transpose, {}],
         [RandomRotate90, {}],
         [Rotate, {}],
+        [SafeRotate, {}],
         [OpticalDistortion, {}],
         [GridDistortion, {}],
         [ElasticTransform, {}],
@@ -569,6 +576,7 @@ def test_image_only_crop_around_bbox_augmentation(augmentation_cls, params, imag
             {"min_height": 514, "min_width": 514, "border_mode": cv2.BORDER_CONSTANT, "value": 100, "mask_value": 1},
         ],
         [Rotate, {"border_mode": cv2.BORDER_CONSTANT, "value": 100, "mask_value": 1}],
+        [SafeRotate, {"border_mode": cv2.BORDER_CONSTANT, "value": 100, "mask_value": 1}],
         [ShiftScaleRotate, {"border_mode": cv2.BORDER_CONSTANT, "value": 100, "mask_value": 1}],
         [OpticalDistortion, {"border_mode": cv2.BORDER_CONSTANT, "value": 100, "mask_value": 1}],
         [ElasticTransform, {"border_mode": cv2.BORDER_CONSTANT, "value": 100, "mask_value": 1}],
@@ -594,6 +602,7 @@ def test_mask_fill_value(augmentation_cls, params):
         [GaussNoise, {}],
         [RandomSizedCrop, {"min_max_height": (384, 512), "height": 512, "width": 512}],
         [ShiftScaleRotate, {}],
+        [SafeRotate, {}],
         [PadIfNeeded, {"min_height": 514, "min_width": 516}],
         [LongestMaxSize, {"max_size": 256}],
         [GridDistortion, {}],
@@ -625,6 +634,7 @@ def test_multichannel_image_augmentations(augmentation_cls, params):
         [GaussNoise, {}],
         [RandomSizedCrop, {"min_max_height": (384, 512), "height": 512, "width": 512}],
         [ShiftScaleRotate, {}],
+        [SafeRotate, {}],
         [PadIfNeeded, {"min_height": 514, "min_width": 516}],
         [LongestMaxSize, {"max_size": 256}],
         [GridDistortion, {}],

tests/test_serialization.py

@@ -54,6 +54,7 @@ def set_seed(seed):
         [A.Transpose, {}],
         [A.RandomRotate90, {}],
         [A.Rotate, {}],
+        [A.SafeRotate, {}],
         [A.ShiftScaleRotate, {}],
         [A.OpticalDistortion, {}],
         [A.GridDistortion, {}],
@@ -167,6 +168,15 @@ def test_augmentations_serialization(augmentation_cls, params, p, seed, image, m
                 "value": (10, 10, 10),
             },
         ],
+        [
+            A.SafeRotate,
+            {
+                "limit": 120,
+                "interpolation": cv2.INTER_CUBIC,
+                "border_mode": cv2.BORDER_CONSTANT,
+                "value": (10, 10, 10),
+            },
+        ],
         [
             A.ShiftScaleRotate,
             {
@@ -349,6 +359,7 @@ def test_augmentations_serialization_to_file_with_custom_parameters(
         [A.Transpose, {}],
         [A.RandomRotate90, {}],
         [A.Rotate, {}],
+        [A.SafeRotate, {}],
         [A.ShiftScaleRotate, {}],
         [A.CenterCrop, {"height": 10, "width": 10}],
         [A.RandomCrop, {"height": 10, "width": 10}],
@@ -423,6 +434,7 @@ def test_augmentations_for_bboxes_serialization(
         [A.Flip, {}],
         [A.RandomRotate90, {}],
         [A.Rotate, {}],
+        [A.SafeRotate, {}],
         [A.ShiftScaleRotate, {}],
         [A.CenterCrop, {"height": 10, "width": 10}],
         [A.RandomCrop, {"height": 10, "width": 10}],
@@ -712,6 +724,7 @@ def test_transform_pipeline_serialization_with_keypoints(seed, image, keypoints,
         [A.Transpose, {}],
         [A.RandomRotate90, {}],
         [A.Rotate, {}],
+        [A.SafeRotate, {}],
         [A.OpticalDistortion, {}],
         [A.GridDistortion, {}],
         [A.ElasticTransform, {}],

tests/test_transforms.py

@@ -27,6 +27,20 @@ def test_transpose_both_image_and_mask():
     assert augmented["mask"].shape == (6, 8)
 
 
+@pytest.mark.parametrize("interpolation", [cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_CUBIC])
+def test_safe_rotate_interpolation(interpolation):
+    image = np.random.randint(low=0, high=256, size=(100, 100, 3), dtype=np.uint8)
+    mask = np.random.randint(low=0, high=2, size=(100, 100), dtype=np.uint8)
+    aug = A.SafeRotate(limit=(45, 45), interpolation=interpolation, p=1)
+    data = aug(image=image, mask=mask)
+    expected_image = FGeometric.safe_rotate(image, 45, interpolation=interpolation, border_mode=cv2.BORDER_REFLECT_101)
+    expected_mask = FGeometric.safe_rotate(
+        mask, 45, interpolation=cv2.INTER_NEAREST, border_mode=cv2.BORDER_REFLECT_101
+    )
+    assert np.array_equal(data["image"], expected_image)
+    assert np.array_equal(data["mask"], expected_mask)
+
+
 @pytest.mark.parametrize("interpolation", [cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_CUBIC])
 def test_rotate_interpolation(interpolation):
     image = np.random.randint(low=0, high=256, size=(100, 100, 3), dtype=np.uint8)
@@ -143,6 +157,7 @@ def test_elastic_transform_interpolation(monkeypatch, interpolation):
         [A.RandomSizedCrop, {"min_max_height": (80, 90), "height": 100, "width": 100}],
         [A.LongestMaxSize, {"max_size": 50}],
         [A.Rotate, {}],
+        [A.SafeRotate, {}],
         [A.OpticalDistortion, {}],
         [A.IAAAffine, {"scale": 1.5}],
         [A.IAAPiecewiseAffine, {"scale": 1.5}],
@@ -170,6 +185,7 @@ def test_binary_mask_interpolation(augmentation_cls, params):
         [A.RandomSizedCrop, {"min_max_height": (80, 90), "height": 100, "width": 100}],
         [A.LongestMaxSize, {"max_size": 50}],
         [A.Rotate, {}],
+        [A.SafeRotate, {}],
         [A.Resize, {"height": 80, "width": 90}],
         [A.Resize, {"height": 120, "width": 130}],
         [A.OpticalDistortion, {}],
@@ -204,6 +220,7 @@ def __test_multiprocessing_support_proc(args):
         [A.RandomSizedCrop, {"min_max_height": (80, 90), "height": 100, "width": 100}],
         [A.LongestMaxSize, {"max_size": 50}],
         [A.Rotate, {}],
+        [A.SafeRotate, {}],
         [A.OpticalDistortion, {}],
         [A.IAAAffine, {"scale": 1.5}],
         [A.IAAPiecewiseAffine, {"scale": 1.5}],
@@ -284,6 +301,7 @@ def test_force_apply():
         [A.Transpose, {}],
         [A.RandomRotate90, {}],
         [A.Rotate, {}],
+        [A.SafeRotate, {}],
         [A.OpticalDistortion, {}],
         [A.GridDistortion, {}],
         [A.ElasticTransform, {}],