[Enhance] Support different border modes in imrotate (#2141)

* support different border_mode in imrotate

* fix docstring

* small fix

* reorder the border_mode arg in imrotate

* fix lint

* polish docstring

* polish docstring

mmcv/image/geometric.py

@@ -1,5 +1,7 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import numbers
+import warnings
+from typing import Optional, Tuple
 
 import cv2
 import numpy as np
@@ -37,6 +39,16 @@ cv2_interp_codes = {
     'lanczos': cv2.INTER_LANCZOS4
 }
 
+cv2_border_modes = {
+    'constant': cv2.BORDER_CONSTANT,
+    'replicate': cv2.BORDER_REPLICATE,
+    'reflect': cv2.BORDER_REFLECT,
+    'wrap': cv2.BORDER_WRAP,
+    'reflect_101': cv2.BORDER_REFLECT_101,
+    'transparent': cv2.BORDER_TRANSPARENT,
+    'isolated': cv2.BORDER_ISOLATED
+}
+
 # Pillow >=v9.1.0 use a slightly different naming scheme for filters.
 # Set pillow_interp_codes according to the naming scheme used.
 if Image is not None:
@@ -301,31 +313,40 @@ def imflip_(img, direction='horizontal'):
         return cv2.flip(img, -1, img)
 
 
-def imrotate(img,
-             angle,
-             center=None,
-             scale=1.0,
-             border_value=0,
-             interpolation='bilinear',
-             auto_bound=False):
+def imrotate(img: np.ndarray,
+             angle: float,
+             center: Optional[Tuple[float, float]] = None,
+             scale: float = 1.0,
+             border_value: int = 0,
+             interpolation: str = 'bilinear',
+             auto_bound: bool = False,
+             border_mode: str = 'constant') -> np.ndarray:
     """Rotate an image.
 
     Args:
-        img (ndarray): Image to be rotated.
+        img (np.ndarray): Image to be rotated.
         angle (float): Rotation angle in degrees, positive values mean
             clockwise rotation.
         center (tuple[float], optional): Center point (w, h) of the rotation in
             the source image. If not specified, the center of the image will be
             used.
         scale (float): Isotropic scale factor.
-        border_value (int): Border value.
+        border_value (int): Border value used in case of a constant border.
+            Defaults to 0.
         interpolation (str): Same as :func:`resize`.
         auto_bound (bool): Whether to adjust the image size to cover the whole
             rotated image.
+        border_mode (str): Pixel extrapolation method. Defaults to 'constant'.
 
     Returns:
-        ndarray: The rotated image.
+        np.ndarray: The rotated image.
     """
+    warnings.warn("We have added an arg 'border_mode' in this func "
+                  'and will reorder the args in the future as: '
+                  '( ..., scale: float = 1.0, '
+                  "border_mode: str = 'constant', "
+                  'border_value: int = 0, ... ). '
+                  'Please use keyword arguments to call this function.')
     if center is not None and auto_bound:
         raise ValueError('`auto_bound` conflicts with `center`')
     h, w = img.shape[:2]
@@ -347,6 +368,7 @@ def imrotate(img,
         img,
         matrix, (w, h),
         flags=cv2_interp_codes[interpolation],
+        borderMode=cv2_border_modes[border_mode],
         borderValue=border_value)
     return rotated
 

tests/test_image/test_geometric.py

@@ -527,6 +527,9 @@ class TestGeometric:
         assert_array_equal(mmcv.imrotate(img, 90, border_value=255), img_r)
         img_r = np.array([[5, 1], [6, 2], [7, 3], [8, 4]])
         assert_array_equal(mmcv.imrotate(img, 90, auto_bound=True), img_r)
+        img_r = np.array([[6, 6, 2, 2], [7, 7, 3, 3]])
+        assert_array_equal(
+            mmcv.imrotate(img, 90, border_mode='replicate'), img_r)
 
         with pytest.raises(ValueError):
             mmcv.imrotate(img, 90, center=(0, 0), auto_bound=True)