You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
79 lines
3.3 KiB
79 lines
3.3 KiB
import os
|
|
import shutil
|
|
import cv2
|
|
import numpy as np
|
|
ROOT_PATH="./train_data/"
|
|
SAVE_PATH='./aug_data/'
|
|
|
|
unified_width=250
|
|
unified_height=250
|
|
#递归删除该目录下的所有文件夹及子文件夹
|
|
if os.path.exists(SAVE_PATH):
|
|
shutil.rmtree(SAVE_PATH)
|
|
#创建多层文件夹，用于存储增强图像数据
|
|
def make_directory(read_path=ROOT_PATH,save_path=SAVE_PATH):
|
|
subfolders = os.listdir(read_path)
|
|
for subfolder in subfolders:
|
|
sub_save_path = save_path+subfolder+'/'
|
|
os.makedirs(sub_save_path)
|
|
|
|
#将文件夹下的所有图片统一某个尺寸
|
|
def resize_images(read_path=ROOT_PATH,save_path=SAVE_PATH,width = unified_width,height = unified_height):
|
|
subfolders = os.listdir(read_path)
|
|
for subfolder in subfolders:
|
|
file_names = os.listdir(read_path+subfolder+'/')
|
|
for filename in file_names:
|
|
path_read_filename = read_path+subfolder+"/"+filename
|
|
img = cv2.imread(path_read_filename)
|
|
img = cv2.resize(img,(width,height),interpolation=cv2.INTER_CUBIC)
|
|
path_save_file = save_path+subfolder+"/"+"resized_"+filename
|
|
cv2.imwrite(path_save_file,img)
|
|
|
|
#直方图均衡化处理
|
|
def equalize_hist_images(read_path = ROOT_PATH,save_path=SAVE_PATH):
|
|
subfolders = os.listdir(read_path)
|
|
for subfolder in subfolders:
|
|
filenames = os.listdir(read_path+subfolder+"/")
|
|
for filename in filenames:
|
|
file_read_path = read_path+subfolder+"/"+filename
|
|
img = cv2.imread(file_read_path)
|
|
b,g,r = cv2.split(img)
|
|
b = cv2.equalizeHist(b)
|
|
g = cv2.equalizeHist(g)
|
|
r = cv2.equalizeHist(r)
|
|
img = cv2.merge([b,g,r])
|
|
file_save_path = save_path+subfolder+"/"+"aug_equ_"+filename
|
|
cv2.imwrite(file_save_path,img)
|
|
|
|
#旋转图片，并扩充保存到指定文件夹下
|
|
def rotate_images(read_path = ROOT_PATH,save_path=SAVE_PATH,num_rotation=3):
|
|
rotation_angle = np.random.randint(-30,30)
|
|
subfolders= os.listdir(read_path)
|
|
for subfolder in subfolders:
|
|
filenames = os.listdir(read_path+subfolder+"/")
|
|
for filename in filenames:
|
|
img = cv2.imread(read_path+subfolder+"/"+filename)
|
|
rows,cols,channel = img.shape
|
|
M = cv2.getRotationMatrix2D(((cols-1)/2.0,(rows-1)/2.0),rotation_angle,1)
|
|
img_rotate = cv2.warpAffine(img,M,(cols,rows))
|
|
rotate_write_path = save_path+subfolder+"/"+"aug_rotate_"+filename
|
|
cv2.imwrite(rotate_write_path,img_rotate)
|
|
|
|
#进行模糊处理
|
|
def blur_images(read_path=ROOT_PATH,save_path=SAVE_PATH):
|
|
blur_size = np.random.randint(0,4)*2+1
|
|
subfolders = os.listdir(read_path)
|
|
for subfolder in subfolders:
|
|
filenames = os.listdir(read_path+subfolder+"/")
|
|
for filename in filenames:
|
|
img = cv2.imread(read_path+subfolder+"/"+filename)
|
|
img_blur = cv2.GaussianBlur(img,(blur_size,blur_size),0) #高斯模糊必须是整数
|
|
blur_file_path = save_path+subfolder+"/"+"Gaussian_"+filename
|
|
cv2.imwrite(blur_file_path,img_blur)
|
|
|
|
if __name__ == '__main__':
|
|
make_directory(ROOT_PATH,SAVE_PATH)
|
|
resize_images()
|
|
rotate_images()
|
|
equalize_hist_images()
|
|
blur_images() |