import os
import cv2
import numpy as np
class DataEnhance:
def __init__(self, train_path, train_save_path, test_path, test_save_path, val_path, val_save_path):
super(DataEnhance, self).__init__()
self.train_path = train_path
self.train_save_path = train_save_path
self.test_path = test_path
self.test_save_path = test_save_path
self.val_path = val_path
self.val_save_path = val_save_path
self.main()
def mirror_image(self, img_path, save_path, i, if_save=True):
"""镜像操作：水平镜像、垂直镜像、水平垂直镜像"""
img = cv2.imread(img_path)
h_flip=cv2.flip(img,1)
v_flip=cv2.flip(img,0)
hv_flip=cv2.flip(img,-1)
if(if_save==True):
cv2.imwrite(save_path + 'h_flip_'+ str(i) + '.png', h_flip)
cv2.imwrite(save_path + 'v_flip' + str(i) + '.png', v_flip)
cv2.imwrite(save_path + 'hv_flip' + str(i) + '.png', hv_flip)
print(save_path + 'h_flip_'+ str(i) + '.png')
print(save_path + 'v_flip' + str(i) + '.png')
print(save_path + 'hv_flip' + str(i) + '.png')
return h_flip,v_flip,hv_flip
def gussian(self, img_path, save_path, i, if_save=True):
"""高斯模糊"""
img = cv2.imread(img_path)
temp = cv2.GaussianBlur(img, (17, 17), 1.5) #高斯模糊
blur_ = cv2.blur(img, (7, 7), (-1, -1)) #均值滤波
if(if_save==True):
cv2.imwrite(save_path + 'gussian_17_' + str(i) +'.png', temp)
cv2.imwrite(save_path + 'blur_7_' + str(i) +'.png', blur_)
print(save_path + 'gussian_17_' + str(i) +'.png')
print(save_path + 'blur_7_' + str(i) +'.png')
return temp, blur_
def addGaussianNoise(self, img_path, save_path, i, percetage=0.3, if_save=True):
"""高斯噪声"""
image = cv2.imread(img_path)
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]
if(if_save==True):
cv2.imwrite(save_path + 'addGaussianNoise_' + str(i) +'.png', G_Noiseimg)
print(save_path + 'addGaussianNoise_' + str(i) +'.png')
def darker(self, img_path, save_path, i, percetage=0.9, if_save=True):
"""变暗"""
image = cv2.imread(img_path)
image_copy = image.copy()
w = image.shape[1]
h = image.shape[0]
#get darker
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)
if(if_save==True):
cv2.imwrite(save_path + 'darker_' + str(i) +'.png', image_copy)
print(save_path + 'darker_' + str(i) +'.png')
return image_copy
def brighter(self, img_path, save_path, i, percetage=1.2, if_save=True):
"""变亮"""
image = cv2.imread(img_path)
image_copy = image.copy()
w = image.shape[1]
h = image.shape[0]
#get brighter
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)
if(if_save==True):
cv2.imwrite(save_path + 'brighter_' + str(i) +'.png', image_copy)
print(save_path + 'brighter_' + str(i) +'.png')
return image_copy
def train_data_Enhance(self, image_path, save_path):
"""
```
对训练集进行数据增强
```
:param image_path: './datasets/train/'
datasets/
train/
Classification_1/
img_1.jpg
img_2.jpg
img_3.jpg
...
Classification_2/
img_1.jpg
img_2.jpg
img_3.jpg
...
test/
...
new_train/
...
new_test/
...
:param save_path: './datasets/new_train/'
:return: None
"""
Classification_list = os.listdir(image_path)
for i in range(len(Classification_list)):
img_path = image_path + Classification_list[i]
self.mirror_image(img_path, save_path, i, if_save=True)
self.gussian(img_path, save_path,i, if_save=True)
self.darker(img_path, save_path, i, percetage=0.9, if_save=True)
self.brighter(img_path, save_path, i, percetage=1.2, if_save=True)
def test_val_data_Enhance(self, image_path, save_path):
"""
```
对测试集和验证集进行数据增强
```
:param image_path: './datasets/train/'
datasets/
train/
Classification_1/
img_1.jpg
img_2.jpg
img_3.jpg
...
Classification_2/
img_1.jpg
img_2.jpg
img_3.jpg
...
test/
...
new_train/
...
new_test/
...
:param save_path: './datasets/new_train/'
:return: None
"""
if not os.path.join(save_path):
os.mkdir(save_path)
self.train_data_Enhance(image_path, save_path)
def GaussianNoise(image, save_path, percetage=0.3):
"""高斯噪声"""
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]
cv2.imwrite(save_path, G_Noiseimg)
Classification_list = os.listdir(image_path)
for i in range(len(Classification_list)):
img_path = image_path + Classification_list[i]
original_ = cv2.imread(img_path)
h_flip, v_flip, hv_flip = self.mirror_image(img_path, save_path, i, if_save=False)
temp, blur_ = self.gussian(img_path, save_path, i, if_save=False)
image_copy_1 = self.darker(img_path, save_path, i, percetage=0.9, if_save=False)
image_copy_2 = self.brighter(img_path, save_path, i, percetage=1.2, if_save=False)
for index, item in enumerate([original_, h_flip, v_flip, hv_flip, temp, blur_, image_copy_1, image_copy_2]):
save = save_path + 'GaussianNoise_new_' + str(i) + str(index) + '.png'
GaussianNoise(item, save, percetage=0.3)
print(save_path + 'GaussianNoise_new_' + str(i) + str(index) + '.png')
def main(self):
self.train_data_Enhance(self.train_path, self.train_save_path)
self.test_val_data_Enhance(self.test_path, self.test_save_path)
self.test_val_data_Enhance(self.val_path, self.val_save_path)
if __name__ == '__main__':
"""图片的输入路径和保存路径"""
train_path = '../datasets/train/pos/'
train_save_path = '../datasets/train_pos/'
test_path = '../datasets/test/pos/'
test_save_path = '../datasets/test_pos/'
val_path = '../datasets/val/pos/'
val_save_path = '../datasets/val_pos/'
DataEnhance(train_path, train_save_path, test_path, test_save_path, val_path, val_save_path)
机器学习图片训练数据增强
猜你喜欢
转载自blog.csdn.net/Stybill_LV_/article/details/110798472
今日推荐
周排行