OpenCV-Python扩充数据集

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

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

此代码需要Python运行环境中包含OpenCV库，安装命令如下：

pip install opencv-python

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
    return img, shape

# 旋转图片
def rotate_image(img):
    imgs = []
    # 旋转90度
    imgs.append(cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE))
    # 旋转180度
    imgs.append(cv2.rotate(img, cv2.ROTATE_180))
    # 旋转270度
    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_x= round(shape[0]/2)
    spilt_y = round(shape[1]/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, 0:spilt_y])
    # 右下
    imgs.append(img[spilt_x:-1, spilt_y:-1])
    # 裁中间
    x = round(shape[0]/5)
    y = round(shape[1]/5)
    imgs.append(img[x:(shape[0]-x), y:(shape[1]-y)])
    return imgs

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

def main():
    # 输入路径
    input_path = r"/home/usr/Desktop/NewFolder/p"
    # 输出路径
    output_path = r"/home/usr/Desktop/NewFolder/p"
    # 计数器，替换文件名
    i = 0

    for img_name in os.listdir(input_path):
        img, shape = 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()