今日做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)
如果有问题欢迎在评论区留言讨论!