mmlab添加albumentations数据增强模块

系列文章目录

mmlab图像增强

---

---

前言

`提示：这里可以添加本文要记录的大概内容：

albumentation是深度学习经典的数据增强模块，本文介绍在mmlab采用albumentation自定义数据增强类

一、mmlab中自定义的数据变换类

该类型需要@TRANSFORMS.register_module()相应的进行修饰，并在跟数据变换的根目录中导入，才能在管道中使用。

二、使用步骤

1.参考代码

官方参考示例：

import random
import mmcv
from mmcv.transforms import BaseTransform, TRANSFORMS
#自定义类需要继承原有的BaseTransform类
@TRANSFORMS.register_module()
class MyFlip(BaseTransform):
    def __init__(self, direction: str):
        super().__init__()
        self.direction = direction
    def transform(self, results: dict) -> dict:
        img = results['img']
        results['img'] = mmcv.imflip(img, direction=self.direction)
        return results

2.使用

代码如下：

import numpy as np

transform = MyFlip(direction='horizontal')
data_dict = {'img': np.random.rand(224, 224, 3)}
data_dict = transform(data_dict)
processed_img = data_dict['img']

三、使用albumentation实现自定义类

1.定义数据变换类

import albumentations as A

@TRANSFORMS.register_module()
class Random_aug(BaseTransform):
    def __init__(self, prob: float):
        super().__init__()
        # self.p = prob
        self.transf = A.Compose([
            # A.GaussNoise(var_limit=(40.0, 50.0),p=1),
            A.RandomBrightness(limit=-0.5,p=prob),
            # A.Downscale (scale_min=0.55, scale_max=0.9,p=1),
            # A.Equalize(p=0.2),
        ])
    def transform(self, results: dict) -> dict:
        img = results['img']
        #需要向image传递数据参数来实现，不然可能会报错
        # self.transf(image = img.astype(np.uint8))
        results['img'] = self.transf(image = img.astype(np.uint8))['image']
        return results

2.使用

import numpy as np

transform1 = Random_aug(prob=0.5)
data_dict = {'img': np.random.rand(224, 224,3)}
data_dict = transform1(results = data_dict)
print(data_dict)
processed_img = data_dict['img']

---

总结

mmengine参考链接