Opencv——数据扩增

今日做yolov5项目时发现数据集中存在部分特征图片比较少，而yolov5内置的数据增强又比较单调，所以用程序对数据进行一步扩增，使一张变成六张

注：文件名称会从1001开始计数，这里可以修改cnt变量改变，该程序仅为一个基础版，所以只扩增图片，不扩增相应的标签文件，如果有需要会再写相应的标签扩增程序（短期内应该是不会写了，看需求吧）。

import os
import random

import numpy as np

import cv2

folder_path = ' ' #想要扩增的图片路径
folder_list = os.listdir(folder_path)
print(folder_list)
print(type(folder_list))
cnt = 1000
def addGaussianNoise(image,percetage,store_path_1,cnt):
    store_path = store_path_1 + str(cnt)+'.jpg' #图片类型默认为jpg，可根据需要修改
    img1 = image.copy()
    w = image.shape[1]
    h = image.shape[0]
    img2 = int(percetage*image.shape[0]*image.shape[1])
    for i in range(img2):
        temp_x = np.random.randint(0,h)
        temp_y = np.random.randint(0,w)
        img1[temp_x][temp_y][np.random.randint(3)] = np.random.randn(1)[0]
    cv2.imwrite(store_path, img1)
def colorjitter(img, cj_type,store_path_1,cnt):

    store_path = store_path_1 + str(cnt)+'.jpg'
    if cj_type == "b":
        # value = random.randint(-50, 50)
        value = np.random.choice(np.array([-50, -40, -30, 30, 40, 50]))
        hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
        h, s, v = cv2.split(hsv)
        if value >= 0:
            lim = 255 - value
            v[v > lim] = 255
            v[v <= lim] += value
        else:
            lim = np.absolute(value)
            v[v < lim] = 0
            v[v >= lim] -= np.absolute(value)

        final_hsv = cv2.merge((h, s, v))
        img1 = cv2.cvtColor(final_hsv, cv2.COLOR_HSV2BGR)
        cv2.imwrite(store_path, img1)

    elif cj_type == "s":
        # value = random.randint(-50, 50)
        value = np.random.choice(np.array([-50, -40, -30, 30, 40, 50]))
        hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
        h, s, v = cv2.split(hsv)
        if value >= 0:
            lim = 255 - value
            s[s > lim] = 255
            s[s <= lim] += value
        else:
            lim = np.absolute(value)
            s[s < lim] = 0
            s[s >= lim] -= np.absolute(value)

        final_hsv = cv2.merge((h, s, v))
        img1 = cv2.cvtColor(final_hsv, cv2.COLOR_HSV2BGR)
        cv2.imwrite(store_path, img1)
    elif cj_type == "c":
        brightness = 10
        contrast = random.randint(40, 100)
        dummy = np.int16(img)
        dummy = dummy * (contrast / 127 + 1) - contrast + brightness
        dummy = np.clip(dummy, 0, 255)
        img1 = np.uint8(dummy)
        cv2.imwrite(store_path, img1)

color=['b','s','c']
for i in folder_list:
    a = folder_path + i
    img = cv2.imread(a)
    store_path_1 = ' ' #这里输入储存的目标文件夹路径
    store_path = store_path_1 + str(cnt)+'.jpg'
    # ————————————————镜像变换——————————————————————
    cnt=cnt+1
    img1 = cv2.flip(img, 1)  # 镜像
    cv2.imwrite(store_path, img1)
for i in folder_list:
    a = folder_path + i
    img = cv2.imread(a)
    store_path_1 = ' ' #这里输入储存的目标文件夹路径
    store_path = store_path_1 + str(cnt) + '.jpg'
    # ————————————————————————平移——————————————————
    cnt=cnt+1
    img_info = img.shape
    height, width = img.shape[:2]
    x = random.uniform(-50, 50)
    y = random.uniform(-50, 50)
    mat_translation = np.float32(
        [[0.8, 0.1, x], [0.1, 0.8, y]])  # 变换矩阵：设置平移变换所需的计算矩阵：2行3列
    img1 = cv2.warpAffine(img, mat_translation, (width, height))  # 变换函数
    cv2.imwrite(store_path, img1)
for i in folder_list:
    a = folder_path + i
    img = cv2.imread(a)
    store_path_1 = ' '#这里输入储存的目标文件夹路径
    store_path = store_path_1 + str(cnt)+'.jpg'
    # ————————————————————饱和度，明暗，色调——————————————————————
    for s in color:
        cnt=cnt+1
        colorjitter(img,s,store_path_1,cnt)
    # ————————————————————————————高斯噪声————————————————————————————
    cnt=cnt+1
    addGaussianNoise(img, 0.3,store_path_1,cnt)

如果有问题欢迎在评论区留言讨论！