add crop_border option to Rotate (#1214)

albumentations/augmentations/geometric/rotate.py

@@ -6,6 +6,7 @@
 import numpy as np
 
 from ...core.transforms_interface import DualTransform, to_tuple
+from ..crops import functional as FCrops
 from . import functional as F
 
 __all__ = ["Rotate", "RandomRotate90", "SafeRotate"]
@@ -63,6 +64,7 @@ class Rotate(DualTransform):
                     list of float): padding value if border_mode is cv2.BORDER_CONSTANT applied for masks.
         method (str): rotation method used for the bounding boxes. Should be one of "largest_box" or "ellipse".
             Default: "largest_box"
+        crop_border (bool): If True would make a largest possible crop within rotated image
         p (float): probability of applying the transform. Default: 0.5.
 
     Targets:
@@ -80,6 +82,7 @@ def __init__(
         value=None,
         mask_value=None,
         method="largest_box",
+        crop_border=False,
         always_apply=False,
         p=0.5,
     ):
@@ -90,27 +93,86 @@ def __init__(
         self.value = value
         self.mask_value = mask_value
         self.method = method
+        self.crop_border = crop_border
 
         if method not in ["largest_box", "ellipse"]:
             raise ValueError(f"Rotation method {self.method} is not valid.")
 
-    def apply(self, img, angle=0, interpolation=cv2.INTER_LINEAR, **params):
-        return F.rotate(img, angle, interpolation, self.border_mode, self.value)
+    def apply(
+        self, img, angle=0, interpolation=cv2.INTER_LINEAR, x_min=None, x_max=None, y_min=None, y_max=None, **params
+    ):
+        img_out = F.rotate(img, angle, interpolation, self.border_mode, self.value)
+        if self.crop_border:
+            img_out = FCrops.crop(img_out, x_min, y_min, x_max, y_max)
+        return img_out
+
+    def apply_to_mask(self, img, angle=0, x_min=None, x_max=None, y_min=None, y_max=None, **params):
+        img_out = F.rotate(img, angle, cv2.INTER_NEAREST, self.border_mode, self.mask_value)
+        if self.crop_border:
+            img_out = FCrops.crop(img_out, x_min, y_min, x_max, y_max)
+        return img_out
+
+    def apply_to_bbox(self, bbox, angle=0, x_min=None, x_max=None, y_min=None, y_max=None, cols=0, rows=0, **params):
+        bbox_out = F.bbox_rotate(bbox, angle, self.method, rows, cols)
+        if self.crop_border:
+            bbox_out = FCrops.bbox_crop(bbox_out, x_min, y_min, x_max, y_max, rows, cols)
+        return bbox_out
+
+    def apply_to_keypoint(
+        self, keypoint, angle=0, x_min=None, x_max=None, y_min=None, y_max=None, cols=0, rows=0, **params
+    ):
+        keypoint_out = F.keypoint_rotate(keypoint, angle, rows, cols, **params)
+        if self.crop_border:
+            keypoint_out = FCrops.crop_keypoint_by_coords(keypoint_out, (x_min, x_max, y_min, y_max))
+        return keypoint_out
 
-    def apply_to_mask(self, img, angle=0, **params):
-        return F.rotate(img, angle, cv2.INTER_NEAREST, self.border_mode, self.mask_value)
+    @staticmethod
+    def _rotated_rect_with_max_area(h, w, angle):
+        """
+        Given a rectangle of size wxh that has been rotated by 'angle' (in
+        degrees), computes the width and height of the largest possible
+        axis-aligned rectangle (maximal area) within the rotated rectangle.
 
-    def get_params(self):
-        return {"angle": random.uniform(self.limit[0], self.limit[1])}
+        Code from: https://stackoverflow.com/questions/16702966/rotate-image-and-crop-out-black-borders
+        """
+
+        angle = math.radians(angle)
+        width_is_longer = w >= h
+        side_long, side_short = (w, h) if width_is_longer else (h, w)
+
+        # since the solutions for angle, -angle and 180-angle are all the same,
+        # it is sufficient to look at the first quadrant and the absolute values of sin,cos:
+        sin_a, cos_a = abs(math.sin(angle)), abs(math.cos(angle))
+        if side_short <= 2.0 * sin_a * cos_a * side_long or abs(sin_a - cos_a) < 1e-10:
+            # half constrained case: two crop corners touch the longer side,
+            # the other two corners are on the mid-line parallel to the longer line
+            x = 0.5 * side_short
+            wr, hr = (x / sin_a, x / cos_a) if width_is_longer else (x / cos_a, x / sin_a)
+        else:
+            # fully constrained case: crop touches all 4 sides
+            cos_2a = cos_a * cos_a - sin_a * sin_a
+            wr, hr = (w * cos_a - h * sin_a) / cos_2a, (h * cos_a - w * sin_a) / cos_2a
+
+        return dict(
+            x_min=max(0, int(w / 2 - wr / 2)),
+            x_max=min(w, int(w / 2 + wr / 2)),
+            y_min=max(0, int(h / 2 - hr / 2)),
+            y_max=min(h, int(h / 2 + hr / 2)),
+        )
 
-    def apply_to_bbox(self, bbox, angle=0, **params):
-        return F.bbox_rotate(bbox, angle, self.method, params["rows"], params["cols"])
+    @property
+    def targets_as_params(self) -> List[str]:
+        return ["image"]
 
-    def apply_to_keypoint(self, keypoint, angle=0, **params):
-        return F.keypoint_rotate(keypoint, angle, **params)
+    def get_params_dependent_on_targets(self, params: Dict[str, Any]) -> Dict[str, Any]:
+        out_params = {"angle": random.uniform(self.limit[0], self.limit[1])}
+        if self.crop_border:
+            h, w = params["image"].shape[:2]
+            out_params.update(self._rotated_rect_with_max_area(h, w, out_params["angle"]))
+        return out_params
 
     def get_transform_init_args_names(self):
-        return ("limit", "interpolation", "border_mode", "value", "mask_value", "method")
+        return ("limit", "interpolation", "border_mode", "value", "mask_value", "method", "crop_border")
 
 
 class SafeRotate(DualTransform):

tests/test_serialization.py

@@ -133,6 +133,7 @@ def test_augmentations_serialization(augmentation_cls, params, p, seed, image, m
             "interpolation": cv2.INTER_CUBIC,
             "border_mode": cv2.BORDER_CONSTANT,
             "value": (10, 10, 10),
+            "crop_border": False,
         },
     ],
     [

tests/test_transforms.py

@@ -38,6 +38,16 @@ def test_rotate_interpolation(interpolation):
     assert np.array_equal(data["mask"], expected_mask)
 
 
+def test_rotate_crop_border():
+    image = np.random.randint(low=100, high=256, size=(100, 100, 3), dtype=np.uint8)
+    border_value = 13
+    aug = A.Rotate(limit=(45, 45), p=1, value=border_value, border_mode=cv2.BORDER_CONSTANT, crop_border=True)
+    aug_img = aug(image=image)["image"]
+    expected_size = int(np.round(100 / np.sqrt(2)))
+    assert aug_img.shape[0] == expected_size
+    assert (aug_img == border_value).sum() == 0
+
+
 @pytest.mark.parametrize("interpolation", [cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_CUBIC])
 def test_shift_scale_rotate_interpolation(interpolation):
     image = np.random.randint(low=0, high=256, size=(100, 100, 3), dtype=np.uint8)