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()