OpenCV-Python扩充数据集

在深度学习中，数据集规模直接影响模型训练后的性能。当我们数据集较少的时候可以通过不同的手段对现有的数据进行扩充。比如裁切，翻转，旋转，加入噪点，调整亮度等手段进行数据集的扩充。

代码演示的是对数据集进行裁切，旋转和镜像操作。代码并不是一键自动化生成，需要根据需求在main()中自行修改。

import cv2
import os

# 加载图片
def load_img(path):
    img=cv2.imread(path)
    if(img is None):
        print("load failed!")
        exit()
    # print(img.shape)
    shape = img.shape
    name = img.name
    return img, shape, name

# 旋转图片
def rotate_image(img):
    imgs = []
    imgs.append(cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE))
    imgs.append(cv2.rotate(img, cv2.ROTATE_180))
    imgs.append(cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE))
    return imgs

# 镜像图片
def mirror_image(img):
    imgs = []
    # 水平镜像
    imgs.append(cv2.flip(img, 1))
    # 垂直镜像
    imgs.append(cv2.flip(img, 0))
    # 对角镜像
    imgs.append(cv2.flip(img, -1))
    return imgs

# 分割图片
def spilt_image(img,shape):
    imgs=[]
    spilt_y = round(shape[1]/2)
    spilt_x= round(shape[0]/2)
    # 左上
    imgs.append(img[0:spilt_x, 0:spilt_y])
    # 左下
    imgs.append(img[0:spilt_x, spilt_y:-1])
    # 右下
    imgs.append(img[spilt_x:-1, spilt_y:-1])
    # 右上
    imgs.append(img[spilt_x:-1, 0:spilt_y])
    return imgs

# 保存图片
def save(img, path):
    cv2.imwrite(path,img)

def main():
    # 输入路径
    input_path = r"/home/shawn.li/Desktop/New Folder/spilt"
    # 输出路径
    output_path = r"/home/shawn.li/Desktop/New Folder/11"
    # 计数器，替换文件名
    i = 0

    for img_name in os.listdir(input_path):
        img, shape, name = load_img(input_path +"/"+ img_name)

        # 分割图片，四份
        imgs = spilt_image(img, shape)
        # 旋转图片，三份
        imgs = rotate_image(img)
        # 镜像图片
        imgs = mirror_image(img)

        for cropped_img in imgs:
            i += 1
            save(cropped_img, output_path+"/"+str(i) +".jpg")
            print("Saved.")


if __name__ == "__main__":
    main()