[Feature] support Albumentations augmentation transformations and pipeline

docs/en/howto/transforms.md

@@ -540,6 +540,47 @@ pipeline = [
    </tbody>
 </table>
 
+### Albumentations
+
+MMagic support adding custom transformations from [Albumentations](https://github.com/albumentations-team/albumentations) library. Please visit https://albumentations.ai/docs/getting_started/transforms_and_targets to get more information.
+
+An example of Albumentations's `transforms` is as followed:
+
+```python
+albu_transforms = [
+   dict(
+         type='Resize',
+         height=100,
+         width=100,
+   ),
+   dict(
+         type='RandomFog',
+         p=0.5,
+   ),
+   dict(
+         type='RandomRain',
+         p=0.5
+   ),
+   dict(
+         type='RandomSnow',
+         p=0.5,
+   ),
+]
+pipeline = [
+   dict(
+         type='LoadImageFromFile',
+         key='img',
+         color_type='color',
+         channel_order='rgb',
+         imdecode_backend='cv2'),
+   dict(
+         type='Albumentations',
+         keys=['img'],
+         transforms=albu_transforms),
+   dict(type='PackInputs')
+]
+```
+
 ## Extend and use custom pipelines
 
 ### A simple example of MyTransform

docs/zh_cn/howto/transforms.md

@@ -540,6 +540,47 @@ pipeline = [
    </tbody>
 </table>
 
+### Albumentations
+
+MMagic 支持添加 [Albumentations](https://github.com/albumentations-team/albumentations) 库中的 transformation，请浏览 https://albumentations.ai/docs/getting_started/transforms_and_targets 获取更多 transformation 的信息。
+
+使用 Albumentations 的示例如下：
+
+```python
+albu_transforms = [
+   dict(
+         type='Resize',
+         height=100,
+         width=100,
+   ),
+   dict(
+         type='RandomFog',
+         p=0.5,
+   ),
+   dict(
+         type='RandomRain',
+         p=0.5
+   ),
+   dict(
+         type='RandomSnow',
+         p=0.5,
+   ),
+]
+pipeline = [
+   dict(
+         type='LoadImageFromFile',
+         key='img',
+         color_type='color',
+         channel_order='rgb',
+         imdecode_backend='cv2'),
+   dict(
+         type='Albumentations',
+         keys=['img'],
+         transforms=albu_transforms),
+   dict(type='PackInputs')
+]
+```
+
 ## 扩展和使用自定义流水线
 
 ### 一个简单的MyTransform示例

mmagic/datasets/transforms/__init__.py

@@ -1,4 +1,5 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+from .albumentations import Albumentations
 from .alpha import GenerateSeg, GenerateSoftSeg
 from .aug_frames import MirrorSequence, TemporalReverse
 from .aug_pixel import (BinarizeImage, Clip, ColorJitter, RandomAffine,
@@ -46,5 +47,6 @@
     'FormatTrimap', 'TransformTrimap', 'GenerateTrimap',
     'GenerateTrimapWithDistTransform', 'CompositeFg', 'RandomLoadResizeBg',
     'MergeFgAndBg', 'PerturbBg', 'RandomJitter', 'LoadPairedImageFromFile',
-    'CenterCropLongEdge', 'RandomCropLongEdge', 'NumpyPad', 'InstanceCrop'
+    'CenterCropLongEdge', 'RandomCropLongEdge', 'NumpyPad', 'InstanceCrop',
+    'Albumentations'
 ]

mmagic/datasets/transforms/albumentations.py

@@ -0,0 +1,142 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import copy
+import inspect
+from typing import List
+
+import numpy as np
+from mmcv.transforms import BaseTransform
+
+from mmagic.registry import TRANSFORMS
+
+try:
+    import albumentations
+    from albumentations import Compose
+except ImportError:
+    albumentations = None
+    Compose = None
+
+
+@TRANSFORMS.register_module()
+class Albumentations(BaseTransform):
+    """Albumentation augmentation.
+
+    Adds custom transformations from Albumentations library.
+    Please, visit `https://github.com/albumentations-team/albumentations`
+    and `https://albumentations.ai/docs/getting_started/transforms_and_targets`
+    to get more information.
+
+    An example of ``transforms`` is as followed:
+
+    .. code-block::
+
+        albu_transforms = [
+            dict(
+                type='Resize',
+                height=100,
+                width=100,
+            ),
+            dict(
+                type='RandomFog',
+                p=0.5,
+            ),
+            dict(
+                type='RandomRain',
+                p=0.5
+            ),
+            dict(
+                type='RandomSnow',
+                p=0.5,
+            ),
+        ]
+        pipeline = [
+            dict(
+                type='LoadImageFromFile',
+                key='img',
+                color_type='color',
+                channel_order='rgb',
+                imdecode_backend='cv2'),
+            dict(
+                type='Albumentations',
+                keys=['img'],
+                transforms=albu_transforms),
+            dict(type='PackInputs')
+        ]
+
+    Args:
+        keys (list[str]): A list specifying the keys whose values are modified.
+        transforms (list[dict]): A list of albu transformations.
+    """
+
+    def __init__(self, keys: List[str], transforms: List[dict]) -> None:
+
+        if Compose is None:
+            raise RuntimeError('Please install albumentations')
+
+        self.keys = keys
+
+        # Args will be modified later, copying it will be safer
+        transforms = copy.deepcopy(transforms)
+        self.transforms = transforms
+        self.aug = Compose([self.albu_builder(t) for t in self.transforms])
+
+    def albu_builder(self, cfg: dict) -> albumentations:
+        """Import a module from albumentations.
+
+        It inherits some of :func:`build_from_cfg` logic.
+
+        Args:
+            cfg (dict): Config dict. It should at least contain the key "type".
+
+        Returns:
+            obj: The constructed object.
+        """
+
+        assert isinstance(cfg, dict) and 'type' in cfg
+        args = cfg.copy()
+        obj_type = args.pop('type')
+        if isinstance(obj_type, str):
+            if albumentations is None:
+                raise RuntimeError('Please install albumentations')
+            obj_cls = getattr(albumentations, obj_type)
+        elif inspect.isclass(obj_type):
+            obj_cls = obj_type
+        else:
+            raise TypeError(
+                f'type must be a str or valid type, but got {type(obj_type)}')
+
+        if 'transforms' in args:
+            args['transforms'] = [
+                self.albu_builder(transform)
+                for transform in args['transforms']
+            ]
+
+        return obj_cls(**args)
+
+    def _apply_albu(self, imgs):
+        is_single_image = False
+        if isinstance(imgs, np.ndarray):
+            is_single_image = True
+            imgs = [imgs]
+
+        outputs = []
+        for img in imgs:
+            outputs.append(self.aug(image=img)['image'])
+
+        if is_single_image:
+            outputs = outputs[0]
+
+        return outputs
+
+    def transform(self, results):
+        """Transform function of Albumentations."""
+
+        for k in self.keys:
+            results[k] = self._apply_albu(results[k])
+
+        return results
+
+    def __repr__(self):
+        repr_str = self.__class__.__name__
+        repr_str += f'(keys={self.keys}, transforms={self.transforms})'
+
+        return repr_str

requirements/optional.txt

@@ -1,3 +1,4 @@
+albumentations
 -e git+https://github.com/openai/CLIP.git@d50d76daa670286dd6cacf3bcd80b5e4823fc8e1#egg=clip
 imageio-ffmpeg==0.4.4
 mmdet >= 3.0.0

requirements/runtime.txt

@@ -3,7 +3,7 @@ av==8.0.3; python_version < '3.7'
 controlnet_aux
 diffusers>=0.12.0
 einops
-face-alignment
+face-alignment<=1.3.4
 facexlib
 lmdb
 lpips

requirements/tests.txt

@@ -1,3 +1,4 @@
+albumentations
 -e git+https://github.com/openai/CLIP.git@d50d76daa670286dd6cacf3bcd80b5e4823fc8e1#egg=clip
 controlnet_aux
 # codecov

tests/test_datasets/test_transforms/test_albumentations.py

@@ -0,0 +1,15 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import numpy as np
+
+from mmagic.datasets.transforms import Albumentations
+
+
+def test_albumentations():
+    results = {}
+    results['img'] = np.ones((8, 8, 3)).astype(np.uint8)
+    model = Albumentations(
+        keys=['img'], transforms=[
+            dict(type='Resize', height=4, width=4),
+        ])
+    results = model(results)
+    assert results['img'].shape == (4, 4, 3)