fix: fix augmentation

Fix RandomErasing, RandAugment to be consistent with Timm and compatible with earlier PaddleClas.
Add ColorJitter implemented by PaddleVision and TimmAutoAugment borrowed from Timm Lib.

ppcls/data/preprocess/__init__.py

@@ -14,6 +14,7 @@
 
 from ppcls.data.preprocess.ops.autoaugment import ImageNetPolicy as RawImageNetPolicy
 from ppcls.data.preprocess.ops.randaugment import RandAugment as RawRandAugment
+from ppcls.data.preprocess.ops.timm_autoaugment import RawTimmAutoAugment
 from ppcls.data.preprocess.ops.cutout import Cutout
 
 from ppcls.data.preprocess.ops.hide_and_seek import HideAndSeek
@@ -31,7 +32,6 @@ from ppcls.data.preprocess.ops.operators import AugMix
 
 from ppcls.data.preprocess.batch_ops.batch_operators import MixupOperator, CutmixOperator, OpSampler, FmixOperator
 
-import six
 import numpy as np
 from PIL import Image
 
@@ -47,20 +47,14 @@ class AutoAugment(RawImageNetPolicy):
     """ ImageNetPolicy wrapper to auto fit different img types """
 
     def __init__(self, *args, **kwargs):
-        if six.PY2:
-            super(AutoAugment, self).__init__(*args, **kwargs)
-        else:
-            super().__init__(*args, **kwargs)
+        super().__init__(*args, **kwargs)
 
     def __call__(self, img):
         if not isinstance(img, Image.Image):
             img = np.ascontiguousarray(img)
             img = Image.fromarray(img)
 
-        if six.PY2:
-            img = super(AutoAugment, self).__call__(img)
-        else:
-            img = super().__call__(img)
+        img = super().__call__(img)
 
         if isinstance(img, Image.Image):
             img = np.asarray(img)
@@ -72,20 +66,33 @@ class RandAugment(RawRandAugment):
     """ RandAugment wrapper to auto fit different img types """
 
     def __init__(self, *args, **kwargs):
-        if six.PY2:
-            super(RandAugment, self).__init__(*args, **kwargs)
-        else:
-            super().__init__(*args, **kwargs)
+        super().__init__(*args, **kwargs)
 
     def __call__(self, img):
         if not isinstance(img, Image.Image):
             img = np.ascontiguousarray(img)
             img = Image.fromarray(img)
 
-        if six.PY2:
-            img = super(RandAugment, self).__call__(img)
-        else:
-            img = super().__call__(img)
+        img = super().__call__(img)
+
+        if isinstance(img, Image.Image):
+            img = np.asarray(img)
+
+        return img
+
+
+class TimmAutoAugment(RawTimmAutoAugment):
+    """ TimmAutoAugment wrapper to auto fit different img tyeps. """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def __call__(self, img):
+        if not isinstance(img, Image.Image):
+            img = np.ascontiguousarray(img)
+            img = Image.fromarray(img)
+
+        img = super().__call__(img)
 
         if isinstance(img, Image.Image):
             img = np.asarray(img)

ppcls/data/preprocess/ops/operators.py

@@ -26,6 +26,7 @@ import random
 import cv2
 import numpy as np
 from PIL import Image
+from paddle.vision.transforms import ColorJitter as RawColorJitter
 
 from .autoaugment import ImageNetPolicy
 from .functional import augmentations
@@ -363,3 +364,20 @@ class AugMix(object):
 
         mixed = (1 - m) * image + m * mix
         return mixed.astype(np.uint8)
+
+
+class ColorJitter(RawColorJitter):
+    """ColorJitter.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def __call__(self, img):
+        if not isinstance(img, Image.Image):
+            img = np.ascontiguousarray(img)
+            img = Image.fromarray(img)
+        img = super()._apply_image(img)
+        if isinstance(img, Image.Image):
+            img = np.asarray(img)
+        return img

ppcls/data/preprocess/ops/random_erasing.py

@@ -12,7 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-#This code is based on https://github.com/zhunzhong07/Random-Erasing
+#This code is adapted from https://github.com/zhunzhong07/Random-Erasing, and refer to Timm.
+
+from functools import partial
 
 import math
 import random
@@ -20,36 +22,69 @@ import random
 import numpy as np
 
 
+class Pixels(object):
+    def __init__(self, mode="const", mean=[0., 0., 0.]):
+        self._mode = mode
+        self._mean = mean
+
+    def __call__(self, h=224, w=224, c=3):
+        if self._mode == "rand":
+            return np.random.normal(size=(1, 1, 3))
+        elif self._mode == "pixel":
+            return np.random.normal(size=(h, w, c))
+        elif self._mode == "const":
+            return self._mean
+        else:
+            raise Exception(
+                "Invalid mode in RandomErasing, only support \"const\", \"rand\", \"pixel\""
+            )
+
+
 class RandomErasing(object):
-    def __init__(self, EPSILON=0.5, sl=0.02, sh=0.4, r1=0.3,
-                 mean=[0., 0., 0.]):
-        self.EPSILON = EPSILON
-        self.mean = mean
-        self.sl = sl
-        self.sh = sh
-        self.r1 = r1
+    """RandomErasing.
+    """
+
+    def __init__(self,
+                 EPSILON=0.5,
+                 sl=0.02,
+                 sh=0.4,
+                 r1=0.3,
+                 mean=[0., 0., 0.],
+                 attempt=100,
+                 use_log_aspect=False,
+                 mode='const'):
+        self.EPSILON = eval(EPSILON) if isinstance(EPSILON, str) else EPSILON
+        self.sl = eval(sl) if isinstance(sl, str) else sl
+        self.sh = eval(sh) if isinstance(sh, str) else sh
+        r1 = eval(r1) if isinstance(r1, str) else r1
+        self.r1 = (math.log(r1), math.log(1 / r1)) if use_log_aspect else (
+            r1, 1 / r1)
+        self.use_log_aspect = use_log_aspect
+        self.attempt = attempt
+        self.get_pixels = Pixels(mode, mean)
 
     def __call__(self, img):
-        if random.uniform(0, 1) > self.EPSILON:
+        if random.random() > self.EPSILON:
             return img
 
-        for _ in range(100):
+        for _ in range(self.attempt):
             area = img.shape[0] * img.shape[1]
 
             target_area = random.uniform(self.sl, self.sh) * area
-            aspect_ratio = random.uniform(self.r1, 1 / self.r1)
+            aspect_ratio = random.uniform(*self.r1)
+            if self.use_log_aspect:
+                aspect_ratio = math.exp(aspect_ratio)
 
             h = int(round(math.sqrt(target_area * aspect_ratio)))
             w = int(round(math.sqrt(target_area / aspect_ratio)))
 
             if w < img.shape[1] and h < img.shape[0]:
+                pixels = self.get_pixels(h, w, img.shape[2])
                 x1 = random.randint(0, img.shape[0] - h)
                 y1 = random.randint(0, img.shape[1] - w)
-                if img.shape[0] == 3:
-                    img[x1:x1 + h, y1:y1 + w, 0] = self.mean[0]
-                    img[x1:x1 + h, y1:y1 + w, 1] = self.mean[1]
-                    img[x1:x1 + h, y1:y1 + w, 2] = self.mean[2]
+                if img.shape[2] == 3:
+                    img[x1:x1 + h, y1:y1 + w, :] = pixels
                 else:
-                    img[0, x1:x1 + h, y1:y1 + w] = self.mean[1]
+                    img[x1:x1 + h, y1:y1 + w, 0] = pixels[0]
                 return img
         return img