【Python-OpenCV】图片数据增广批量转换模板

import cv2
import numpy as np
import os.path


# 图像增加椒盐噪声函数
def SaltAndPepper(src, percetage):
    SP_NoiseImg = src.copy()
    SP_NoiseNum = int(percetage * src.shape[0] * src.shape[1])
    for i in range(SP_NoiseNum):
        randR = np.random.randint(0, src.shape[0] - 1)
        randG = np.random.randint(0, src.shape[1] - 1)
        randB = np.random.randint(0, 3)
        if np.random.randint(0, 1) == 0:
            SP_NoiseImg[randR, randG, randB] = 0
        else:
            SP_NoiseImg[randR, randG, randB] = 255
    return SP_NoiseImg


# 图像增加高斯噪声函数
def AddGaussianNoise(image, percetage):
    G_Noiseimg = image.copy()
    w = image.shape[1]
    h = image.shape[0]
    G_NoiseNum = int(percetage * image.shape[0] * image.shape[1])
    for i in range(G_NoiseNum):
        temp_x = np.random.randint(0, h)
        temp_y = np.random.randint(0, w)
        G_Noiseimg[temp_x][temp_y][np.random.randint(3)] = np.random.randn(1)[0]
    return G_Noiseimg


# 图像暗化处理函数
def Darker(image, percetage=0.9):
    image_copy = image.copy()
    w = image.shape[1]
    h = image.shape[0]
    for xi in range(0, w):
        for xj in range(0, h):
            image_copy[xj, xi, 0] = int(image[xj, xi, 0] * percetage)
            image_copy[xj, xi, 1] = int(image[xj, xi, 1] * percetage)
            image_copy[xj, xi, 2] = int(image[xj, xi, 2] * percetage)
    return image_copy


# 图像亮化处理函数
def Brighter(image, percetage=1.5):
    image_copy = image.copy()
    w = image.shape[1]
    h = image.shape[0]
    for xi in range(0, w):
        for xj in range(0, h):
            image_copy[xj, xi, 0] = np.clip(int(image[xj, xi, 0] * percetage), a_max=255, a_min=0)
            image_copy[xj, xi, 1] = np.clip(int(image[xj, xi, 1] * percetage), a_max=255, a_min=0)
            image_copy[xj, xi, 2] = np.clip(int(image[xj, xi, 2] * percetage), a_max=255, a_min=0)
    return image_copy


# 图像翻转随机角度函数
def Rotate(image, angle, center=None, scale=1.0):
    (h, w) = image.shape[:2]
    if center is None:
        center = (w / 2, h / 2)
    m = cv2.getRotationMatrix2D(center, angle, scale)
    rotated = cv2.warpAffine(image, m, (w, h))
    return rotated


# 图像镜像翻转函数
def Flip(image):
    flipped_image = np.fliplr(image)
    return flipped_image


file_dir = "img/"  # 同级图像文件夹，根据自己的文件目录路径来设置

# 图像数据增广模板（可根据实际需要注释或修改参数）
for img_name in os.listdir(file_dir):

    # 读入单幅图像
    img_path = file_dir + img_name
    img = cv2.imread(img_path)

    # 翻转90°和180°
    rotated_90 = Rotate(img, 90)  # 度数可自己设定
    cv2.imwrite(file_dir + img_name[0:-4] + '_r90.jpg', rotated_90)    # 翻转90°
    rotated_180 = Rotate(img, 180)  # 度数可自己设定
    cv2.imwrite(file_dir + img_name[0:-4] + '_r180.jpg', rotated_180)  # 翻转180°

    # 镜像翻转
    flipped_img = Flip(img)
    cv2.imwrite(file_dir + img_name[0:-4] + '_fli.jpg', flipped_img)

    # 增加噪声（参数可自己设定）
    img_salt = SaltAndPepper(img, 0.3)      # 椒盐噪声
    cv2.imwrite(file_dir + img_name[0:7] + '_salt.jpg', img_salt)
    img_gauss = AddGaussianNoise(img, 0.3)  # 高斯噪声
    cv2.imwrite(file_dir + img_name[0:-4] + '_noise.jpg', img_gauss)

    # 变亮变暗（参数可自己设定）
    img_darker = Darker(img)      # 变暗
    cv2.imwrite(file_dir + img_name[0:-4] + '_darker.jpg', img_darker)
    img_brighter = Brighter(img)  # 变亮
    cv2.imwrite(file_dir + img_name[0:-4] + '_brighter.jpg', img_brighter)

    # 高斯滤波（参数可自己设定）
    blur = cv2.GaussianBlur(img, (7, 7), 1.5)  # 参数：图像、卷积核、标准差
    cv2.imwrite(file_dir + img_name[0:-4] + '_blur.jpg', blur)

print("图像数据增强完毕！")