Add elliptical rotation option to RotationTransform

gluoncv/torch/data/transforms/instance_transforms/transform.py

@@ -968,13 +968,14 @@ def inverse(self):
         return ResizeTransform(self.new_h, self.new_w, self.h, self.w, self.interp)
 
 
+
 class RotationTransform(Transform):
     """
     This method returns a copy of this image, rotated the given
     number of degrees counter clockwise around its center.
     """
 
-    def __init__(self, h, w, angle, expand=True, center=None, interp=None):
+    def __init__(self, h, w, angle, expand=True, center=None, interp=None, box_method="largest"):
         """
         Args:
             h, w (int): original image size
@@ -985,6 +986,11 @@ def __init__(self, h, w, angle, expand=True, center=None, interp=None):
                 if left to None, the center will be fit to the center of each image
                 center has no effect if expand=True because it only affects shifting
             interp: cv2 interpolation method, default cv2.INTER_LINEAR
+            box_method: either `'largest'` (default) or `'ellipse'`. Affects how bboxes are rotated.
+                If `'ellipse'`, then bboxes are rotated as if the object is an ellipse rather than a
+                rectangle, which avoids creating oversized bounding boxes. (see https://arxiv.org/abs/2109.13488)
+                If `'largest'`, this will rotate the corner points and use their minimum/maximum
+                to create a new axis-aligned box.
         """
         super().__init__()
         image_center = np.array((w / 2, h / 2))
@@ -1006,6 +1012,10 @@ def __init__(self, h, w, angle, expand=True, center=None, interp=None):
         # Needed because of this problem https://github.com/opencv/opencv/issues/11784
         self.rm_image = self.create_rotation_matrix(offset=-0.5)
 
+        if box_method not in ["largest", "ellipse"]:
+            raise ValueError(f"Method '{box_method}' is not a valid box rotation method.")
+        self.box_method = box_method
+
     def apply_image(self, img, interp=None):
         """
         img should be a numpy array, formatted as Height * Width * Nchannels
@@ -1055,6 +1065,31 @@ def inverse(self):
         )
         return TransformList([rotation, crop])
 
+    def apply_box(self, box):
+        """
+        box should be a Nx4 floating point array of XYXY format in absolute coordinates.
+        """
+        if self.box_method == "largest":
+            return super().apply_box(box)
+        else:
+            box = np.asarray(box)
+            w, h = box[:, 2] - box[:, 0], box[:, 3] - box[:, 1]
+            x, y = box[:, 0] + w / 2, box[:, 1] + h / 2
+
+            # Create 32 keypoints along ellipsis
+            coords = [
+                [x + np.sin(t) * (w / 2), y + np.cos(t) * (h / 2)]
+                for t in np.arange(0, 2 * np.pi, 0.2)
+            ]  # 32x2xN
+            coords = np.moveaxis(coords, 2, 0).reshape(-1, 2)  # (N*32)x2
+
+            # Transform these coordinates in the same way as rectangle coordinates
+            coords = self.apply_coords(coords).reshape(-1, 32, 2)  # Nx32x2
+            minxy = coords.min(axis=1)
+            maxxy = coords.max(axis=1)
+            trans_boxes = np.concatenate((minxy, maxxy), axis=1)
+            return trans_boxes
+
 
 def HFlip_rotated_box(transform, rotated_boxes):
     """