数据增强——在图片中添加遮挡物

致歉

之前没有好好检查代码就把代码发上来了，导致博客上的代码无法运行，向大家道歉，现在已更新代码

为了增强网络的分类泛化能力，准备在图片中加入一些遮挡物。代码如下：

import random

import itertools

import os
import PIL.Image as Image
import PIL.ImageDraw as ImageDraw

# 原始图片的存放位置
PATH = './source_path/'
# 新生成的图片的保存位置
SAVE_PATH= './save_path/'
# 要在图片上生成几边形的物体，N=5代表五边形
N = 5


def drawObs(path, savePath, n):

    for file in os.listdir(path):
        if not file.lower().endswith(('.bmp', '.dib', '.png', '.jpg', '.jpeg', '.pbm', '.pgm', '.ppm', '.tif', '.tiff')):
            continue
        img = Image.open(path + file)
        (x, y) = img.size
        random_list = list(itertools.product(range(1, x), range(1, y)))
        points = random.sample(random_list, n)

        x1 = random.randint(0, 255)
        x2 = random.randint(0, 255)
        x3 = random.randint(0, 255)

        draw = ImageDraw.Draw(img)

        ImageDraw.ImageDraw.polygon(xy=points, fill=(x1, x2, x3), self=draw)
        img.save(savePath + 'rand' + file)
        print(file)

if __name__ == '__main__':
    drawObs(PATH, SAVE_PATH, N)

随机在图片的某个位置加入五边形的遮挡物，颜色随机。效果如下：

处理前

处理后