样本扩充完整代码超级实用 python

柔情只为你懂 2022-02-28 10:24 466阅读 0赞

这是博主用来做样本扩充常用的代码，里面备注已经很详细了，直接拿来用吧！

其中autoaugment.py直接借用的github的代码，很棒。

### 一.代码 ###

1.  autoaugment.py

'自动数据增强方法AutoAugment'
    """
    基于扩充样本操作的流程：比如先找来500张车辆的图片，并标注得xml文件，再转化得txt文件。
    然后用此autoaugment.py做扩充操作—>比如每张原图扩充为10张，那么这10张对应的新txt文件里的坐标还是同一个！
    这样的好处是扩充后只是加了各种噪声、锐化、光暗之类的，不需要再对新生成的图做xml坐标标注了！
    故：为了不需再对新生成的图做xml坐标标注，我把这个autoaugment.py的类ImageNetPolicy(object)里
    所有引用到了类SubPolicy(object)里的旋转操作rotate的代码全屏蔽掉了。
    """
    from PIL import Image, ImageEnhance, ImageOps
    import numpy as np
    import random
    
    class ImageNetPolicy(object):
        """ Randomly choose one of the best 24 Sub-policies on ImageNet.
    
            Example:
            # >>> policy = ImageNetPolicy()
            # >>> transformed = policy(image)
    
            Example as a PyTorch Transform:
            # >>> transform=transforms.Compose([
            # >>>     transforms.Resize(256),
            # >>>     ImageNetPolicy(),
            # >>>     transforms.ToTensor()])
        """
        def __init__(self, fillcolor=(128, 128, 128)):
            self.policies = [
                # SubPolicy(0.4, "posterize", 8, 0.6, "rotate", 9, fillcolor),
                SubPolicy(0.6, "solarize", 5, 0.6, "autocontrast", 5, fillcolor),
                SubPolicy(0.8, "equalize", 8, 0.6, "equalize", 3, fillcolor),
                SubPolicy(0.6, "posterize", 7, 0.6, "posterize", 6, fillcolor),
                SubPolicy(0.4, "equalize", 7, 0.2, "solarize", 4, fillcolor),
    
                # SubPolicy(0.4, "equalize", 4, 0.8, "rotate", 8, fillcolor),
                SubPolicy(0.6, "solarize", 3, 0.6, "equalize", 7, fillcolor),
                SubPolicy(0.8, "posterize", 5, 1.0, "equalize", 2, fillcolor),
                # SubPolicy(0.2, "rotate", 3, 0.6, "solarize", 8, fillcolor),
                SubPolicy(0.6, "equalize", 8, 0.4, "posterize", 6, fillcolor),
    
                # SubPolicy(0.8, "rotate", 8, 0.4, "color", 0, fillcolor),
                # SubPolicy(0.4, "rotate", 9, 0.6, "equalize", 2, fillcolor),
                SubPolicy(0.0, "equalize", 7, 0.8, "equalize", 8, fillcolor),
                SubPolicy(0.6, "invert", 4, 1.0, "equalize", 8, fillcolor),
                SubPolicy(0.6, "color", 4, 1.0, "contrast", 8, fillcolor),
    
                # SubPolicy(0.8, "rotate", 8, 1.0, "color", 2, fillcolor),
                SubPolicy(0.8, "color", 8, 0.8, "solarize", 7, fillcolor),
                SubPolicy(0.4, "sharpness", 7, 0.6, "invert", 8, fillcolor),
                # SubPolicy(0.6, "shearX", 5, 1.0, "equalize", 9, fillcolor),
                SubPolicy(0.4, "color", 0, 0.6, "equalize", 3, fillcolor),
    
                SubPolicy(0.4, "equalize", 7, 0.2, "solarize", 4, fillcolor),
                SubPolicy(0.6, "solarize", 5, 0.6, "autocontrast", 5, fillcolor),
                SubPolicy(0.6, "invert", 4, 1.0, "equalize", 8, fillcolor),
                SubPolicy(0.6, "color", 4, 1.0, "contrast", 8, fillcolor)
            ]
    
        def __call__(self, img):
            policy_idx = random.randint(0, len(self.policies) - 1)
            print(policy_idx)
            return self.policies[policy_idx](img)
    
        def __repr__(self):
            return "AutoAugment ImageNet Policy"
    
    class CIFAR10Policy(object):
        """ Randomly choose one of the best 25 Sub-policies on CIFAR10.
    
            Example:
            # >>> policy = CIFAR10Policy()
            # >>> transformed = policy(image)
            #
            Example as a PyTorch Transform:
            # >>> transform=transforms.Compose([
            # >>>     transforms.Resize(256),
            # >>>     CIFAR10Policy(),
            # >>>     transforms.ToTensor()])
        """
        def __init__(self, fillcolor=(128, 128, 128)):
            self.policies = [
                SubPolicy(0.1, "invert", 7, 0.2, "contrast", 6, fillcolor),
                SubPolicy(0.7, "rotate", 2, 0.3, "translateX", 9, fillcolor),
                SubPolicy(0.8, "sharpness", 1, 0.9, "sharpness", 3, fillcolor),
                SubPolicy(0.5, "shearY", 8, 0.7, "translateY", 9, fillcolor),
                SubPolicy(0.5, "autocontrast", 8, 0.9, "equalize", 2, fillcolor),
    
                SubPolicy(0.2, "shearY", 7, 0.3, "posterize", 7, fillcolor),
                SubPolicy(0.4, "color", 3, 0.6, "brightness", 7, fillcolor),
                SubPolicy(0.3, "sharpness", 9, 0.7, "brightness", 9, fillcolor),
                SubPolicy(0.6, "equalize", 5, 0.5, "equalize", 1, fillcolor),
                SubPolicy(0.6, "contrast", 7, 0.6, "sharpness", 5, fillcolor),
    
                SubPolicy(0.7, "color", 7, 0.5, "translateX", 8, fillcolor),
                SubPolicy(0.3, "equalize", 7, 0.4, "autocontrast", 8, fillcolor),
                SubPolicy(0.4, "translateY", 3, 0.2, "sharpness", 6, fillcolor),
                SubPolicy(0.9, "brightness", 6, 0.2, "color", 8, fillcolor),
                SubPolicy(0.5, "solarize", 2, 0.0, "invert", 3, fillcolor),
    
                SubPolicy(0.2, "equalize", 0, 0.6, "autocontrast", 0, fillcolor),
                SubPolicy(0.2, "equalize", 8, 0.8, "equalize", 4, fillcolor),
                SubPolicy(0.9, "color", 9, 0.6, "equalize", 6, fillcolor),
                SubPolicy(0.8, "autocontrast", 4, 0.2, "solarize", 8, fillcolor),
                SubPolicy(0.1, "brightness", 3, 0.7, "color", 0, fillcolor),
    
                SubPolicy(0.4, "solarize", 5, 0.9, "autocontrast", 3, fillcolor),
                SubPolicy(0.9, "translateY", 9, 0.7, "translateY", 9, fillcolor),
                SubPolicy(0.9, "autocontrast", 2, 0.8, "solarize", 3, fillcolor),
                SubPolicy(0.8, "equalize", 8, 0.1, "invert", 3, fillcolor),
                SubPolicy(0.7, "translateY", 9, 0.9, "autocontrast", 1, fillcolor)
            ]
    
        def __call__(self, img):
            policy_idx = random.randint(0, len(self.policies) - 1)
            return self.policies[policy_idx](img)
    
        def __repr__(self):
            return "AutoAugment CIFAR10 Policy"
    
    
    class SVHNPolicy(object):
        """ Randomly choose one of the best 25 Sub-policies on SVHN.
    
            Example:
            # >>> policy = SVHNPolicy()
            # >>> transformed = policy(image)
    
            Example as a PyTorch Transform:
            # >>> transform=transforms.Compose([
            # >>>     transforms.Resize(256),
            # >>>     SVHNPolicy(),
            # >>>     transforms.ToTensor()])
        """
        def __init__(self, fillcolor=(128, 128, 128)):
            self.policies = [
                SubPolicy(0.9, "shearX", 4, 0.2, "invert", 3, fillcolor),
                SubPolicy(0.9, "shearY", 8, 0.7, "invert", 5, fillcolor),
                SubPolicy(0.6, "equalize", 5, 0.6, "solarize", 6, fillcolor),
                SubPolicy(0.9, "invert", 3, 0.6, "equalize", 3, fillcolor),
                SubPolicy(0.6, "equalize", 1, 0.9, "rotate", 3, fillcolor),
    
                SubPolicy(0.9, "shearX", 4, 0.8, "autocontrast", 3, fillcolor),
                SubPolicy(0.9, "shearY", 8, 0.4, "invert", 5, fillcolor),
                SubPolicy(0.9, "shearY", 5, 0.2, "solarize", 6, fillcolor),
                SubPolicy(0.9, "invert", 6, 0.8, "autocontrast", 1, fillcolor),
                SubPolicy(0.6, "equalize", 3, 0.9, "rotate", 3, fillcolor),
    
                SubPolicy(0.9, "shearX", 4, 0.3, "solarize", 3, fillcolor),
                SubPolicy(0.8, "shearY", 8, 0.7, "invert", 4, fillcolor),
                SubPolicy(0.9, "equalize", 5, 0.6, "translateY", 6, fillcolor),
                SubPolicy(0.9, "invert", 4, 0.6, "equalize", 7, fillcolor),
                SubPolicy(0.3, "contrast", 3, 0.8, "rotate", 4, fillcolor),
    
                SubPolicy(0.8, "invert", 5, 0.0, "translateY", 2, fillcolor),
                SubPolicy(0.7, "shearY", 6, 0.4, "solarize", 8, fillcolor),
                SubPolicy(0.6, "invert", 4, 0.8, "rotate", 4, fillcolor),
                SubPolicy(0.3, "shearY", 7, 0.9, "translateX", 3, fillcolor),
                SubPolicy(0.1, "shearX", 6, 0.6, "invert", 5, fillcolor),
    
                SubPolicy(0.7, "solarize", 2, 0.6, "translateY", 7, fillcolor),
                SubPolicy(0.8, "shearY", 4, 0.8, "invert", 8, fillcolor),
                SubPolicy(0.7, "shearX", 9, 0.8, "translateY", 3, fillcolor),
                SubPolicy(0.8, "shearY", 5, 0.7, "autocontrast", 3, fillcolor),
                SubPolicy(0.7, "shearX", 2, 0.1, "invert", 5, fillcolor)
            ]
    
        def __call__(self, img):
            policy_idx = random.randint(0, len(self.policies) - 1)
            return self.policies[policy_idx](img)
    
        def __repr__(self):
            return "AutoAugment SVHN Policy"
    
    
    class SubPolicy(object):
        def __init__(self, p1, operation1, magnitude_idx1, p2, operation2, magnitude_idx2, fillcolor=(128, 128, 128)):
            ranges = {
                "shearX": np.linspace(0, 0.3, 10),
                "shearY": np.linspace(0, 0.3, 10),
                "translateX": np.linspace(0, 150 / 331, 10),
                "translateY": np.linspace(0, 150 / 331, 10),
                "rotate": np.linspace(0, 30, 10),
                "color": np.linspace(0.0, 0.9, 10),
                "posterize": np.round(np.linspace(8, 4, 10), 0).astype(np.int),
                "solarize": np.linspace(256, 0, 10),
                "contrast": np.linspace(0.0, 0.9, 10),
                "sharpness": np.linspace(0.0, 0.9, 10),
                "brightness": np.linspace(0.0, 0.9, 10),
                "autocontrast": [0] * 10,
                "equalize": [0] * 10,
                "invert": [0] * 10
            }
    
            # from https://stackoverflow.com/questions/5252170/specify-image-filling
            # -color-when-rotating-in-python-with-pil-and-setting-expand
            def rotate_with_fill(img, magnitude):
                rot = img.convert("RGBA").rotate(magnitude)
                return Image.composite(rot, Image.new("RGBA", rot.size, (128,) * 4), rot).convert(img.mode)
    
            func = {
                "shearX": lambda img, magnitude: img.transform(
                    img.size, Image.AFFINE, (1, magnitude * random.choice([-1, 1]), 0, 0, 1, 0),
                    Image.BICUBIC, fillcolor=fillcolor),
                "shearY": lambda img, magnitude: img.transform(
                    img.size, Image.AFFINE, (1, 0, 0, magnitude * random.choice([-1, 1]), 1, 0),
                    Image.BICUBIC, fillcolor=fillcolor),
                "translateX": lambda img, magnitude: img.transform(
                    img.size, Image.AFFINE, (1, 0, magnitude * img.size[0] * random.choice([-1, 1]), 0, 1, 0),
                    fillcolor=fillcolor),
                "translateY": lambda img, magnitude: img.transform(
                    img.size, Image.AFFINE, (1, 0, 0, 0, 1, magnitude * img.size[1] * random.choice([-1, 1])),
                    fillcolor=fillcolor),
                "rotate": lambda img, magnitude: rotate_with_fill(img, magnitude),
                # "rotate": lambda img, magnitude: img.rotate(magnitude * random.choice([-1, 1])),
                "color": lambda img, magnitude: ImageEnhance.Color(img).enhance(1 + magnitude * random.choice([-1, 1])),
                "posterize": lambda img, magnitude: ImageOps.posterize(img, magnitude),
                "solarize": lambda img, magnitude: ImageOps.solarize(img, magnitude),
                "contrast": lambda img, magnitude: ImageEnhance.Contrast(img).enhance(
                    1 + magnitude * random.choice([-1, 1])),
                "sharpness": lambda img, magnitude: ImageEnhance.Sharpness(img).enhance(
                    1 + magnitude * random.choice([-1, 1])),
                "brightness": lambda img, magnitude: ImageEnhance.Brightness(img).enhance(
                    1 + magnitude * random.choice([-1, 1])),
                "autocontrast": lambda img, magnitude: ImageOps.autocontrast(img),
                "equalize": lambda img, magnitude: ImageOps.equalize(img),
                "invert": lambda img, magnitude: ImageOps.invert(img)
            }
    
            # self.name = "{}_{:.2f}_and_{}_{:.2f}".format(
            #     operation1, ranges[operation1][magnitude_idx1],
            #     operation2, ranges[operation2][magnitude_idx2])
            self.p1 = p1
            self.operation1 = func[operation1]
            self.magnitude1 = ranges[operation1][magnitude_idx1]
            self.p2 = p2
            self.operation2 = func[operation2]
            self.magnitude2 = ranges[operation2][magnitude_idx2]
    
        def __call__(self, img):
            if random.random() < self.p1:
                img = self.operation1(img, self.magnitude1)
            if random.random() < self.p2:
                img = self.operation2(img, self.magnitude2)
            return img

2.    generate.py

from autoaugment import ImageNetPolicy
    import PIL
    
    for i in range(1):
        '读'
        path = r"C:\Users\admin\Desktop\t\img/"
        f = open(r"C:\Users\admin\Desktop\t\xml.txt", 'r')
        '写'
        save_dir = r"C:\Users\admin\Desktop\t\img_/"
        f1 = open(r"C:\Users\admin\Desktop\t\xml_.txt", 'a')
    
        j = 0
        for lines in f.readlines():
            try:
                for k in range(10):
                    line = lines.strip().split(',',1)
                    # print(line)
                    # print(line[0])
                    img = PIL.Image.open(path + line[0])#读图
                    policy = ImageNetPolicy()
                    img1 = policy(img) #扩充
    
                    f1.write(r"{}.jpg".format(j) +','+line[1]+'\n') #生成txt
                    img1.save(save_dir + '{}.jpg'.format(j)) #生成图
                    j += 1
    
            except Exception as e:
                pass #当某处出错，直接跳过

### 二.效果展示 ###

文件夹img、xml、xml.txt是原本的， img\_是存放扩充后图片的，xml\_.txt是对应img\_的：

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM5OTM4NjY2_size_16_color_FFFFFF_t_70][]

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM5OTM4NjY2_size_16_color_FFFFFF_t_70 1][]

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM5OTM4NjY2_size_16_color_FFFFFF_t_70 2][]

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM5OTM4NjY2_size_16_color_FFFFFF_t_70 3][]

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM5OTM4NjY2_size_16_color_FFFFFF_t_70 4][]

[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM5OTM4NjY2_size_16_color_FFFFFF_t_70]: /images/20220228/508fe07710e148a3bba170dd6fbcbfc5.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM5OTM4NjY2_size_16_color_FFFFFF_t_70 1]: /images/20220228/7630803468114e2ea082e71e12ce8cf2.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM5OTM4NjY2_size_16_color_FFFFFF_t_70 2]: /images/20220228/cc7c0bd46c2c4195ac3e275bf113e69c.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM5OTM4NjY2_size_16_color_FFFFFF_t_70 3]: /images/20220228/11ffa98a8edc4be49f92a82c9198bcc8.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM5OTM4NjY2_size_16_color_FFFFFF_t_70 4]: /images/20220228/14f5acff54b64a9d9db73f6ff113b9e5.png