[Feature] Add cutout and lighting (#909)

* add cutout

* add adjust_lighting

* minor fix

mmcv/image/__init__.py

@@ -2,15 +2,16 @@
 from .colorspace import (bgr2gray, bgr2hls, bgr2hsv, bgr2rgb, bgr2ycbcr,
                          gray2bgr, gray2rgb, hls2bgr, hsv2bgr, imconvert,
                          rgb2bgr, rgb2gray, rgb2ycbcr, ycbcr2bgr, ycbcr2rgb)
-from .geometric import (imcrop, imflip, imflip_, impad, impad_to_multiple,
-                        imrescale, imresize, imresize_like, imrotate, imshear,
-                        imtranslate, rescale_size)
+from .geometric import (cutout, imcrop, imflip, imflip_, impad,
+                        impad_to_multiple, imrescale, imresize, imresize_like,
+                        imrotate, imshear, imtranslate, rescale_size)
 from .io import imfrombytes, imread, imwrite, supported_backends, use_backend
 from .misc import tensor2imgs
 from .photometric import (adjust_brightness, adjust_color, adjust_contrast,
-                          adjust_sharpness, auto_contrast, clahe,
-                          imdenormalize, imequalize, iminvert, imnormalize,
-                          imnormalize_, lut_transform, posterize, solarize)
+                          adjust_lighting, adjust_sharpness, auto_contrast,
+                          clahe, imdenormalize, imequalize, iminvert,
+                          imnormalize, imnormalize_, lut_transform, posterize,
+                          solarize)
 
 __all__ = [
     'bgr2gray', 'bgr2hls', 'bgr2hsv', 'bgr2rgb', 'gray2bgr', 'gray2rgb',
@@ -22,5 +23,5 @@ __all__ = [
     'rgb2ycbcr', 'bgr2ycbcr', 'ycbcr2rgb', 'ycbcr2bgr', 'tensor2imgs',
     'imshear', 'imtranslate', 'adjust_color', 'imequalize',
     'adjust_brightness', 'adjust_contrast', 'lut_transform', 'clahe',
-    'adjust_sharpness', 'auto_contrast'
+    'adjust_sharpness', 'auto_contrast', 'cutout', 'adjust_lighting'
 ]

mmcv/image/geometric.py

@@ -468,6 +468,61 @@ def impad_to_multiple(img, divisor, pad_val=0):
     return impad(img, shape=(pad_h, pad_w), pad_val=pad_val)
 
 
+def cutout(img, shape, pad_val=0):
+    """Randomly cut out a rectangle from the original img.
+
+    Args:
+        img (ndarray): Image to be cutout.
+        shape (int | tuple[int]): Expected cutout shape (h, w). If given as a
+            int, the value will be used for both h and w.
+        pad_val (int | float | tuple[int | float]): Values to be filled in the
+            cut area. Defaults to 0.
+
+    Returns:
+        ndarray: The cutout image.
+    """
+
+    channels = 1 if img.ndim == 2 else img.shape[2]
+    if isinstance(shape, int):
+        cut_h, cut_w = shape, shape
+    else:
+        assert isinstance(shape, tuple) and len(shape) == 2, \
+            f'shape must be a int or a tuple with length 2, but got type ' \
+            f'{type(shape)} instead.'
+        cut_h, cut_w = shape
+    if isinstance(pad_val, (int, float)):
+        pad_val = tuple([pad_val] * channels)
+    elif isinstance(pad_val, tuple):
+        assert len(pad_val) == channels, \
+            'Expected the num of elements in tuple equals the channels' \
+            'of input image. Found {} vs {}'.format(
+                len(pad_val), channels)
+    else:
+        raise TypeError(f'Invalid type {type(pad_val)} for `pad_val`')
+
+    img_h, img_w = img.shape[:2]
+    y0 = np.random.uniform(img_h)
+    x0 = np.random.uniform(img_w)
+
+    y1 = int(max(0, y0 - cut_h / 2.))
+    x1 = int(max(0, x0 - cut_w / 2.))
+    y2 = min(img_h, y1 + cut_h)
+    x2 = min(img_w, x1 + cut_w)
+
+    if img.ndim == 2:
+        patch_shape = (y2 - y1, x2 - x1)
+    else:
+        patch_shape = (y2 - y1, x2 - x1, channels)
+
+    img_cutout = img.copy()
+    patch = np.array(
+        pad_val, dtype=img.dtype) * np.ones(
+            patch_shape, dtype=img.dtype)
+    img_cutout[y1:y2, x1:x2, ...] = patch
+
+    return img_cutout
+
+
 def _get_shear_matrix(magnitude, direction='horizontal'):
     """Generate the shear matrix for transformation.
 

mmcv/image/photometric.py

@@ -332,6 +332,49 @@ def adjust_sharpness(img, factor=1., kernel=None):
     return sharpened_img.astype(img.dtype)
 
 
+def adjust_lighting(img, eigval, eigvec, alphastd=0.1, to_rgb=True):
+    """AlexNet-style PCA jitter.
+
+    This data augmentation is proposed in `ImageNet Classification with Deep
+    Convolutional Neural Networks
+    <https://dl.acm.org/doi/pdf/10.1145/3065386>`_.
+
+    Args:
+        img (ndarray): Image to be ajusted lighting. BGR order.
+        eigval (ndarray): the eigenvalue of the convariance matrix of pixel
+            values, respectively.
+        eigvec (ndarray): the eigenvector of the convariance matrix of pixel
+            values, respectively.
+        alphastd (float): The standard deviation for distribution of alpha.
+            Dafaults to 0.1
+        to_rgb (bool): Whether to convert img to rgb.
+
+    Returns:
+        ndarray: The adjusted image.
+    """
+    assert isinstance(eigval, np.ndarray) and isinstance(eigvec, np.ndarray), \
+        f'eigval and eigvec should both be of type np.ndarray, got ' \
+        f'{type(eigval)} and {type(eigvec)} instead.'
+
+    assert eigval.ndim == 1 and eigvec.ndim == 2
+    assert eigvec.shape == (3, eigval.shape[0])
+    n_eigval = eigval.shape[0]
+    assert isinstance(alphastd, float), 'alphastd should be of type float, ' \
+        f'got {type(alphastd)} instead.'
+
+    img = img.copy().astype(np.float32)
+    if to_rgb:
+        cv2.cvtColor(img, cv2.COLOR_BGR2RGB, img)  # inplace
+
+    alpha = np.random.normal(0, alphastd, n_eigval)
+    alter = eigvec \
+        * np.broadcast_to(alpha.reshape(1, n_eigval), (3, n_eigval)) \
+        * np.broadcast_to(eigval.reshape(1, n_eigval), (3, n_eigval))
+    alter = np.broadcast_to(alter.sum(axis=1).reshape(1, 1, 3), img.shape)
+    img_adjusted = img + alter
+    return img_adjusted
+
+
 def lut_transform(img, lut_table):
     """Transform array by look-up table.
 

tests/test_image/test_geometric.py

@@ -423,6 +423,40 @@ class TestGeometric:
         padded_img = mmcv.impad_to_multiple(img, 2)
         assert padded_img.shape == (20, 12)
 
+    def test_cutout(self):
+        img = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]).astype(np.uint8)
+
+        # shape must be int or tuple
+        with pytest.raises(AssertionError):
+            mmcv.cutout(img, 2.5)
+        # pad_val must be int or float or tuple with the same length
+        # of img channels
+        with pytest.raises(AssertionError):
+            mmcv.cutout(img, 1, (1, 2, 3))
+        with pytest.raises(TypeError):
+            mmcv.cutout(img, 1, None)
+
+        # test cutout the whole img
+        assert_array_equal(mmcv.cutout(img, 6), np.zeros_like(img))
+        # test not cutout
+        assert_array_equal(mmcv.cutout(img, 0), img)
+        # test cutout when shape is int
+        np.random.seed(0)
+        img_cutout = np.array([[1, 2, 3], [4, 0, 6], [7, 8,
+                                                      9]]).astype(np.uint8)
+        assert_array_equal(mmcv.cutout(img, 1), img_cutout)
+        img_cutout = np.array([[1, 2, 3], [4, 10, 6], [7, 8,
+                                                       9]]).astype(np.uint8)
+        assert_array_equal(mmcv.cutout(img, 1, pad_val=10), img_cutout)
+        # test cutout when shape is tuple
+        np.random.seed(0)
+        img_cutout = np.array([[1, 2, 3], [0, 0, 6], [7, 8,
+                                                      9]]).astype(np.uint8)
+        assert_array_equal(mmcv.cutout(img, (1, 2)), img_cutout)
+        img_cutout = np.array([[1, 2, 3], [10, 10, 6], [7, 8,
+                                                        9]]).astype(np.uint8)
+        assert_array_equal(mmcv.cutout(img, (1, 2), pad_val=10), img_cutout)
+
     def test_imrotate(self):
         img = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]).astype(np.uint8)
         assert_array_equal(mmcv.imrotate(img, 0), img)

tests/test_image/test_photometric.py

@@ -291,6 +291,28 @@ class TestPhotometric:
                 rtol=0,
                 atol=1)
 
+    def test_adjust_lighting(self):
+        img = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]).astype(np.uint8)
+        img = np.stack([img, img, img], axis=-1)
+
+        # eigval and eigvec must be np.ndarray
+        with pytest.raises(AssertionError):
+            mmcv.adjust_lighting(img, 1, np.ones((3, 1)))
+        with pytest.raises(AssertionError):
+            mmcv.adjust_lighting(img, np.array([1]), (1, 1, 1))
+        # we must have the same number of eigval and eigvec
+        with pytest.raises(AssertionError):
+            mmcv.adjust_lighting(img, np.array([1]), np.eye(2))
+        with pytest.raises(AssertionError):
+            mmcv.adjust_lighting(img, np.array([1]), np.array([1]))
+
+        img_adjusted = mmcv.adjust_lighting(
+            img,
+            np.random.normal(0, 1, 2),
+            np.random.normal(0, 1, (3, 2)),
+            alphastd=0.)
+        assert_array_equal(img_adjusted, img)
+
     def test_lut_transform(self):
         lut_table = np.array(list(range(256)))