Fix safe rotate targets (#1109)

* Fix safe rotate targets

* Save all keypoints and bboxes

* Fix worng angle error

* Accurate safe rotate

* Tests

* Fix mask interpolation, tests, remove unused code

* Remove dummy test

* Fix mypy

albumentations/augmentations/geometric/functional.py

@@ -44,7 +44,6 @@
     "safe_rotate",
     "bbox_safe_rotate",
     "keypoint_safe_rotate",
-    "safe_rotate_enlarged_img_size",
     "piecewise_affine",
     "to_distance_maps",
     "from_distance_maps",
@@ -539,112 +538,75 @@ def bbox_affine(
 @preserve_channel_dim
 def safe_rotate(
     img: np.ndarray,
-    angle: int = 0,
-    interpolation: int = cv2.INTER_LINEAR,
-    value: int = None,
+    matrix: np.ndarray,
+    interpolation: int,
+    value: Optional[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(
+) -> np.ndarray:
+    h, w = img.shape[:2]
+    warp_fn = _maybe_process_in_chunks(
         cv2.warpAffine,
-        M=rotation_mat,
-        dsize=(new_cols, new_rows),
+        M=matrix,
+        dsize=(w, h),
         flags=interpolation,
         borderMode=border_mode,
         borderValue=value,
     )
+    return warp_fn(img)
 
-    # 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,
+def bbox_safe_rotate(
+    bbox: Tuple[float, float, float, float], matrix: np.ndarray, cols: int, rows: int
+) -> Tuple[float, float, float, float]:
+    x1, y1, x2, y2 = denormalize_bbox(bbox, rows, cols)
+    points = np.array(
+        [
+            [x1, y1, 1],
+            [x2, y1, 1],
+            [x2, y2, 1],
+            [x1, y2, 1],
+        ]
     )
+    points = points @ matrix.T
+    x1 = points[:, 0].min()
+    x2 = points[:, 0].max()
+    y1 = points[:, 1].min()
+    y2 = points[:, 1].max()
 
-    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])
+    def fix_point(pt1: float, pt2: float, max_val: float) -> Tuple[float, float]:
+        # In my opinion, these errors should be very low, around 1-2 pixels.
+        if pt1 < 0:
+            return 0, pt2 + pt1
+        if pt2 > max_val:
+            return pt1 - (pt2 - max_val), max_val
+        return pt1, pt2
 
-    # Rotate keypoint
-    rotated_keypoint = keypoint_rotate(shifted_keypoint, angle, rows=new_rows, cols=new_cols)
+    x1, x2 = fix_point(x1, x2, cols)
+    y1, y2 = fix_point(y1, y2, rows)
 
-    # Scale the keypoint
-    return keypoint_scale(rotated_keypoint, old_cols / new_cols, old_rows / new_rows)
+    return normalize_bbox((x1, y1, x2, y2), rows, cols)
 
 
-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)
+def keypoint_safe_rotate(
+    keypoint: Tuple[float, float, float, float],
+    matrix: np.ndarray,
+    angle: float,
+    scale_x: float,
+    scale_y: float,
+    cols: int,
+    rows: int,
+) -> Tuple[float, float, float, float]:
+    x, y, a, s = keypoint
+    point = np.array([[x, y, 1]])
+    x, y = (point @ matrix.T)[0]
 
-    # The height of the frame to contain the rotated image
-    r_rows = cols * np.sin(angle) + rows * np.cos(angle)
+    # To avoid problems with float errors
+    x = np.clip(x, 0, cols - 1)
+    y = np.clip(y, 0, rows - 1)
 
-    # 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)
+    a += angle
+    s *= max(scale_x, scale_y)
+    return x, y, a, s
 
 
 @clipped

albumentations/augmentations/geometric/rotate.py

@@ -1,4 +1,6 @@
+import math
 import random
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
 
 import cv2
 import numpy as np
@@ -135,13 +137,13 @@ class SafeRotate(DualTransform):
 
     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,
+        limit: Union[float, Tuple[float, float]] = 90,
+        interpolation: int = cv2.INTER_LINEAR,
+        border_mode: int = cv2.BORDER_REFLECT_101,
+        value: Optional[Union[int, float, Sequence[int], Sequence[float]]] = None,
+        mask_value: Optional[Union[int, float, Sequence[int], Sequence[float]]] = None,
+        always_apply: bool = False,
+        p: float = 0.5,
     ):
         super(SafeRotate, self).__init__(always_apply, p)
         self.limit = to_tuple(limit)
@@ -150,24 +152,70 @@ def __init__(
         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(self, img: np.ndarray, matrix: np.ndarray = None, **params) -> np.ndarray:
+        return F.safe_rotate(img, matrix, self.interpolation, self.value, 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 apply_to_mask(self, img: np.ndarray, matrix: np.ndarray = None, **params) -> np.ndarray:
+        return F.safe_rotate(img, matrix, cv2.INTER_NEAREST, self.mask_value, 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_bbox(
+        self, bbox: Tuple[float, float, float, float], cols: int = 0, rows: int = 0, **params
+    ) -> Tuple[float, float, float, float]:
+        return F.bbox_safe_rotate(bbox, params["matrix"], cols, rows)
 
-    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):
+    def apply_to_keypoint(
+        self,
+        keypoint: Tuple[float, float, float, float],
+        angle: float = 0,
+        scale_x: float = 0,
+        scale_y: float = 0,
+        cols: int = 0,
+        rows: int = 0,
+        **params
+    ) -> Tuple[float, float, float, float]:
+        return F.keypoint_safe_rotate(keypoint, params["matrix"], angle, scale_x, scale_y, cols, rows)
+
+    @property
+    def targets_as_params(self) -> List[str]:
+        return ["image"]
+
+    def get_params_dependent_on_targets(self, params: Dict[str, Any]) -> Dict[str, Any]:
+        angle = random.uniform(self.limit[0], self.limit[1])
+
+        image = params["image"]
+        h, w = image.shape[:2]
+
+        # https://stackoverflow.com/questions/43892506/opencv-python-rotate-image-without-cropping-sides
+        image_center = (w / 2, h / 2)
+
+        # Rotation Matrix
+        rotation_mat = cv2.getRotationMatrix2D(image_center, angle, 1.0)
+
+        # rotation calculates the cos and sin, taking absolutes of those.
+        abs_cos = abs(rotation_mat[0, 0])
+        abs_sin = abs(rotation_mat[0, 1])
+
+        # find the new width and height bounds
+        new_w = math.ceil(h * abs_sin + w * abs_cos)
+        new_h = math.ceil(h * abs_cos + w * abs_sin)
+
+        scale_x = w / new_w
+        scale_y = h / new_h
+
+        # Shift the image to create padding
+        rotation_mat[0, 2] += new_w / 2 - image_center[0]
+        rotation_mat[1, 2] += new_h / 2 - image_center[1]
+
+        # Rescale to original size
+        scale_mat = np.diag(np.ones(3))
+        scale_mat[0, 0] *= scale_x
+        scale_mat[1, 1] *= scale_y
+        _tmp = np.diag(np.ones(3))
+        _tmp[:2] = rotation_mat
+        _tmp = scale_mat @ _tmp
+        rotation_mat = _tmp[:2]
+
+        return {"matrix": rotation_mat, "angle": angle, "scale_x": scale_x, "scale_y": scale_y}
+
+    def get_transform_init_args_names(self) -> Tuple[str, str, str, str, str]:
         return ("limit", "interpolation", "border_mode", "value", "mask_value")

albumentations/augmentations/keypoints_utils.py

@@ -160,11 +160,11 @@ def convert_keypoint_from_albumentations(
     # type (tuple, str, int, int, bool, bool) -> tuple
     if target_format not in keypoint_formats:
         raise ValueError("Unknown target_format {}. Supported formats are: {}".format(target_format, keypoint_formats))
-    if check_validity:
-        check_keypoint(keypoint, rows, cols)
 
     (x, y, angle, scale), tail = keypoint[:4], tuple(keypoint[4:])
     angle = angle_to_2pi_range(angle)
+    if check_validity:
+        check_keypoint((x, y, angle, scale), rows, cols)
     if angle_in_degrees:
         angle = math.degrees(angle)