diff --git a/mmcv/image/geometric.py b/mmcv/image/geometric.py index eecd795ea..066f53988 100644 --- a/mmcv/image/geometric.py +++ b/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 diff --git a/tests/test_image/test_geometric.py b/tests/test_image/test_geometric.py index 308f96c1c..e6409d7e5 100644 --- a/tests/test_image/test_geometric.py +++ b/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)