[Feature] Support AutoAug, AutoContrast, Equalize, Contrast, Brightne…

…ss and Sharpness (#179)

* add AutoContrast, Equalize, Contrast, Brightness and Sharpness pipelines

* add ImageNetPolicy

* add configs

* add unittest

* remove config

* rerun CI

* rerun CI

* [Fix] Update pip install mmcv command in ci (#187)

* update pip install mmcv command in ci

* update pip install mmcv command in ci

* fix ci

* fix ci

mmcls/datasets/pipelines/__init__.py

@@ -1,5 +1,7 @@
-from .auto_augment import (ColorTransform, Invert, Posterize, Rotate, Shear,
-                           Solarize, Translate)
+from .auto_augment import (AutoAugment, AutoContrast, Brightness,
+                           ColorTransform, Contrast, Equalize, Invert,
+                           Posterize, Rotate, Sharpness, Shear, Solarize,
+                           Translate)
 from .compose import Compose
 from .formating import (Collect, ImageToTensor, ToNumpy, ToPIL, ToTensor,
                         Transpose, to_tensor)
@@ -12,5 +14,6 @@
     'Transpose', 'Collect', 'LoadImageFromFile', 'Resize', 'CenterCrop',
     'RandomFlip', 'Normalize', 'RandomCrop', 'RandomResizedCrop',
     'RandomGrayscale', 'Shear', 'Translate', 'Rotate', 'Invert',
-    'ColorTransform', 'Solarize', 'Posterize'
+    'ColorTransform', 'Solarize', 'Posterize', 'AutoContrast', 'Equalize',
+    'Contrast', 'Brightness', 'Sharpness', 'AutoAugment'
 ]

mmcls/datasets/pipelines/auto_augment.py

@@ -1,14 +1,55 @@
+import copy
+
 import mmcv
 import numpy as np
 
 from ..builder import PIPELINES
+from .compose import Compose
 
 
 def random_negative(value, random_negative_prob):
     """Randomly negate value based on random_negative_prob."""
     return -value if np.random.rand() < random_negative_prob else value
 
 
+@PIPELINES.register_module()
+class AutoAugment(object):
+    """Auto augmentation.
+    This data augmentation is proposed in `AutoAugment: Learning Augmentation
+    Policies from Data <https://arxiv.org/abs/1805.09501>`_.
+
+    Args:
+        policies (list[list[dict]]): The policies of auto augmentation. Each
+            policy in ``policies`` is a specific augmentation policy, and is
+            composed by several augmentations (dict). When AutoAugment is
+            called, a random policy in ``policies`` will be selected to
+            augment images.
+    """
+
+    def __init__(self, policies):
+        assert isinstance(policies, list) and len(policies) > 0, \
+            'Policies must be a non-empty list.'
+        for policy in policies:
+            assert isinstance(policy, list) and len(policy) > 0, \
+                'Each policy in policies must be a non-empty list.'
+            for augment in policy:
+                assert isinstance(augment, dict) and 'type' in augment, \
+                    'Each specific augmentation must be a dict with key' \
+                    ' "type".'
+
+        self.policies = copy.deepcopy(policies)
+        self.sub_policy = [Compose(policy) for policy in self.policies]
+
+    def __call__(self, results):
+        sub_policy = np.random.choice(self.sub_policy)
+        return sub_policy(results)
+
+    def __repr__(self):
+        repr_str = self.__class__.__name__
+        repr_str += f'(policies={self.policies})'
+        return repr_str
+
+
 @PIPELINES.register_module()
 class Shear(object):
     """Shear images.
@@ -261,6 +302,36 @@ def __repr__(self):
         return repr_str
 
 
+@PIPELINES.register_module()
+class AutoContrast(object):
+    """Auto adjust image contrast.
+
+    Args:
+        prob (float): The probability for performing invert therefore should
+             be in range [0, 1]. Defaults to 0.5.
+    """
+
+    def __init__(self, prob=0.5):
+        assert 0 <= prob <= 1.0, 'The prob should be in range [0,1], ' \
+            f'got {prob} instead.'
+
+        self.prob = prob
+
+    def __call__(self, results):
+        if np.random.rand() > self.prob:
+            return results
+        for key in results.get('img_fields', ['img']):
+            img = results[key]
+            img_contrasted = mmcv.auto_contrast(img)
+            results[key] = img_contrasted.astype(img.dtype)
+        return results
+
+    def __repr__(self):
+        repr_str = self.__class__.__name__
+        repr_str += f'(prob={self.prob})'
+        return repr_str
+
+
 @PIPELINES.register_module()
 class Invert(object):
     """Invert images.
@@ -292,52 +363,39 @@ def __repr__(self):
 
 
 @PIPELINES.register_module()
-class ColorTransform(object):
-    """Adjust the color balance of images.
+class Equalize(object):
+    """Equalize the image histogram.
 
     Args:
-        magnitude (int | float): The magnitude used for color transform. A
-            positive magnitude would enhance the color and a negative magnitude
-             would make the image grayer. A magnitude=0 gives the origin img.
-        prob (float): The probability for performing ColorTransform therefore
-            should be in range [0, 1]. Defaults to 0.5.
-        random_negative_prob (float): The probability that turns the magnitude
-            negative, which should be in range [0,1]. Defaults to 0.5.
+        prob (float): The probability for performing invert therefore should
+             be in range [0, 1]. Defaults to 0.5.
     """
 
-    def __init__(self, magnitude, prob=0.5, random_negative_prob=0.5):
-        assert isinstance(magnitude, (int, float)), 'The magnitude type must '\
-            f'be int or float, but got {type(magnitude)} instead.'
+    def __init__(self, prob=0.5):
         assert 0 <= prob <= 1.0, 'The prob should be in range [0,1], ' \
             f'got {prob} instead.'
-        assert 0 <= random_negative_prob <= 1.0, 'The random_negative_prob ' \
-            f'should be in range [0,1], got {random_negative_prob} instead.'
 
-        self.magnitude = magnitude
         self.prob = prob
-        self.random_negative_prob = random_negative_prob
 
     def __call__(self, results):
         if np.random.rand() > self.prob:
             return results
-        magnitude = random_negative(self.magnitude, self.random_negative_prob)
         for key in results.get('img_fields', ['img']):
             img = results[key]
-            img_color_adjusted = mmcv.adjust_color(img, alpha=1 + magnitude)
-            results[key] = img_color_adjusted.astype(img.dtype)
+            img_equalized = mmcv.imequalize(img)
+            results[key] = img_equalized.astype(img.dtype)
         return results
 
     def __repr__(self):
         repr_str = self.__class__.__name__
-        repr_str += f'(magnitude={self.magnitude}, '
-        repr_str += f'prob={self.prob}, '
-        repr_str += f'random_negative_prob={self.random_negative_prob})'
+        repr_str += f'(prob={self.prob})'
         return repr_str
 
 
 @PIPELINES.register_module()
 class Solarize(object):
-    """Solarize an image (invert all pixel values above a threshold).
+    """Solarize images (invert all pixel values above a threshold).
+
     Args:
         thr (int | float): The threshold above which the pixels value will be
             inverted.
@@ -372,7 +430,8 @@ def __repr__(self):
 
 @PIPELINES.register_module()
 class Posterize(object):
-    """Posterize an image (reduce the number of bits for each color channel).
+    """Posterize images (reduce the number of bits for each color channel).
+
     Args:
         bits (int): Number of bits for each pixel in the output img, which
             should be less or equal to 8.
@@ -404,3 +463,182 @@ def __repr__(self):
         repr_str += f'(bits={self.bits}, '
         repr_str += f'prob={self.prob})'
         return repr_str
+
+
+@PIPELINES.register_module()
+class Contrast(object):
+    """Adjust images contrast.
+
+    Args:
+        magnitude (int | float): The magnitude used for adjusting contrast. A
+            positive magnitude would enhance the contrast and a negative
+            magnitude would make the image grayer. A magnitude=0 gives the
+            origin img.
+        prob (float): The probability for performing contrast adjusting
+            therefore should be in range [0, 1]. Defaults to 0.5.
+        random_negative_prob (float): The probability that turns the magnitude
+            negative, which should be in range [0,1]. Defaults to 0.5.
+    """
+
+    def __init__(self, magnitude, prob=0.5, random_negative_prob=0.5):
+        assert isinstance(magnitude, (int, float)), 'The magnitude type must '\
+            f'be int or float, but got {type(magnitude)} instead.'
+        assert 0 <= prob <= 1.0, 'The prob should be in range [0,1], ' \
+            f'got {prob} instead.'
+        assert 0 <= random_negative_prob <= 1.0, 'The random_negative_prob ' \
+            f'should be in range [0,1], got {random_negative_prob} instead.'
+
+        self.magnitude = magnitude
+        self.prob = prob
+        self.random_negative_prob = random_negative_prob
+
+    def __call__(self, results):
+        if np.random.rand() > self.prob:
+            return results
+        magnitude = random_negative(self.magnitude, self.random_negative_prob)
+        for key in results.get('img_fields', ['img']):
+            img = results[key]
+            img_contrasted = mmcv.adjust_contrast(img, factor=1 + magnitude)
+            results[key] = img_contrasted.astype(img.dtype)
+        return results
+
+    def __repr__(self):
+        repr_str = self.__class__.__name__
+        repr_str += f'(magnitude={self.magnitude}, '
+        repr_str += f'prob={self.prob}, '
+        repr_str += f'random_negative_prob={self.random_negative_prob})'
+        return repr_str
+
+
+@PIPELINES.register_module()
+class ColorTransform(object):
+    """Adjust images color balance.
+
+    Args:
+        magnitude (int | float): The magnitude used for color transform. A
+            positive magnitude would enhance the color and a negative magnitude
+             would make the image grayer. A magnitude=0 gives the origin img.
+        prob (float): The probability for performing ColorTransform therefore
+            should be in range [0, 1]. Defaults to 0.5.
+        random_negative_prob (float): The probability that turns the magnitude
+            negative, which should be in range [0,1]. Defaults to 0.5.
+    """
+
+    def __init__(self, magnitude, prob=0.5, random_negative_prob=0.5):
+        assert isinstance(magnitude, (int, float)), 'The magnitude type must '\
+            f'be int or float, but got {type(magnitude)} instead.'
+        assert 0 <= prob <= 1.0, 'The prob should be in range [0,1], ' \
+            f'got {prob} instead.'
+        assert 0 <= random_negative_prob <= 1.0, 'The random_negative_prob ' \
+            f'should be in range [0,1], got {random_negative_prob} instead.'
+
+        self.magnitude = magnitude
+        self.prob = prob
+        self.random_negative_prob = random_negative_prob
+
+    def __call__(self, results):
+        if np.random.rand() > self.prob:
+            return results
+        magnitude = random_negative(self.magnitude, self.random_negative_prob)
+        for key in results.get('img_fields', ['img']):
+            img = results[key]
+            img_color_adjusted = mmcv.adjust_color(img, alpha=1 + magnitude)
+            results[key] = img_color_adjusted.astype(img.dtype)
+        return results
+
+    def __repr__(self):
+        repr_str = self.__class__.__name__
+        repr_str += f'(magnitude={self.magnitude}, '
+        repr_str += f'prob={self.prob}, '
+        repr_str += f'random_negative_prob={self.random_negative_prob})'
+        return repr_str
+
+
+@PIPELINES.register_module()
+class Brightness(object):
+    """Adjust images brightness.
+
+    Args:
+        magnitude (int | float): The magnitude used for adjusting brightness. A
+            positive magnitude would enhance the brightness and a negative
+            magnitude would make the image darker. A magnitude=0 gives the
+            origin img.
+        prob (float): The probability for performing contrast adjusting
+            therefore should be in range [0, 1]. Defaults to 0.5.
+        random_negative_prob (float): The probability that turns the magnitude
+            negative, which should be in range [0,1]. Defaults to 0.5.
+    """
+
+    def __init__(self, magnitude, prob=0.5, random_negative_prob=0.5):
+        assert isinstance(magnitude, (int, float)), 'The magnitude type must '\
+            f'be int or float, but got {type(magnitude)} instead.'
+        assert 0 <= prob <= 1.0, 'The prob should be in range [0,1], ' \
+            f'got {prob} instead.'
+        assert 0 <= random_negative_prob <= 1.0, 'The random_negative_prob ' \
+            f'should be in range [0,1], got {random_negative_prob} instead.'
+
+        self.magnitude = magnitude
+        self.prob = prob
+        self.random_negative_prob = random_negative_prob
+
+    def __call__(self, results):
+        if np.random.rand() > self.prob:
+            return results
+        magnitude = random_negative(self.magnitude, self.random_negative_prob)
+        for key in results.get('img_fields', ['img']):
+            img = results[key]
+            img_brightened = mmcv.adjust_brightness(img, factor=1 + magnitude)
+            results[key] = img_brightened.astype(img.dtype)
+        return results
+
+    def __repr__(self):
+        repr_str = self.__class__.__name__
+        repr_str += f'(magnitude={self.magnitude}, '
+        repr_str += f'prob={self.prob}, '
+        repr_str += f'random_negative_prob={self.random_negative_prob})'
+        return repr_str
+
+
+@PIPELINES.register_module()
+class Sharpness(object):
+    """Adjust images sharpness.
+
+    Args:
+        magnitude (int | float): The magnitude used for adjusting sharpness. A
+            positive magnitude would enhance the sharpness and a negative
+            magnitude would make the image bulr. A magnitude=0 gives the
+            origin img.
+        prob (float): The probability for performing contrast adjusting
+            therefore should be in range [0, 1]. Defaults to 0.5.
+        random_negative_prob (float): The probability that turns the magnitude
+            negative, which should be in range [0,1]. Defaults to 0.5.
+    """
+
+    def __init__(self, magnitude, prob=0.5, random_negative_prob=0.5):
+        assert isinstance(magnitude, (int, float)), 'The magnitude type must '\
+            f'be int or float, but got {type(magnitude)} instead.'
+        assert 0 <= prob <= 1.0, 'The prob should be in range [0,1], ' \
+            f'got {prob} instead.'
+        assert 0 <= random_negative_prob <= 1.0, 'The random_negative_prob ' \
+            f'should be in range [0,1], got {random_negative_prob} instead.'
+
+        self.magnitude = magnitude
+        self.prob = prob
+        self.random_negative_prob = random_negative_prob
+
+    def __call__(self, results):
+        if np.random.rand() > self.prob:
+            return results
+        magnitude = random_negative(self.magnitude, self.random_negative_prob)
+        for key in results.get('img_fields', ['img']):
+            img = results[key]
+            img_sharpened = mmcv.adjust_sharpness(img, factor=1 + magnitude)
+            results[key] = img_sharpened.astype(img.dtype)
+        return results
+
+    def __repr__(self):
+        repr_str = self.__class__.__name__
+        repr_str += f'(magnitude={self.magnitude}, '
+        repr_str += f'prob={self.prob}, '
+        repr_str += f'random_negative_prob={self.random_negative_prob})'
+        return repr_str