'''
日期:2020/7/8  16:44
修改人：鲁易
'''
# import torch
import torchvision.transforms.functional as TF
from torchvision import transforms
from PIL import Image
import os

def __ImageTrans(x):    #目前只放了一个变换（非label的图像才会用）
  im_aug = transforms.Compose([
    transforms.ColorJitter(brightness=0.3, contrast=0.3, hue=0.3),
  ])
  x = im_aug(x)
  return x

def rotate_and_save(image_dir, target_dir, angles_list, islabel):
    """
    指定一些角度旋转输入图像文件夹并保存
    :param image_dir:输入图像的路径目录（文件夹）
    :param target_dir:目标图片存储目录（文件夹）
    :param angles_list: 需要对输入图像进行多少角度的旋转
    :param islabel:是否是label图，如果是的话就只旋转，不做其他变换
    """
    if not os.path.exists(target_dir):  #文件夹不存在就新建
        os.mkdir(target_dir)
    for root, dirs, files in os.walk(image_dir):
        for file in files:
            image_path = root + '/' + file
            image = Image.open(image_path)
            # 创建旋转后的图像容器
            for angle in angles_list:
                rotated_image = TF.rotate(img=image, angle=angle, resample=Image.NEAREST)
                if (islabel == False):  # 不是label的话，就不止旋转了
                    rotated_image = __ImageTrans(rotated_image);
                # 按照原始image_path的图片名称作为目标图像名称
                # 获取图片名称并去掉".JPG"后缀
                rotated_image.save(
                    os.path.join(target_dir, os.path.basename(image_path))[:-4] + "_" + str(angle) + "_angel" + "_.jpg")

if __name__ == '__main__':
    angle_list = []
    for i in range(0,361,30):
        angle_list.append(i)
    print(angle_list)
    #下面就是一个简单的例子，需要最好把每个数据集的文件都这么操作一下
    image_dir, target_dir = r"F:\PROGRAM\SHM\ICP-SHM-P1\Dataset for Training\2-80 Additional Original Images",\
                            r"F:\PROGRAM\SHM\ICP-SHM-P1\Dataset for Training\2-80 Additional Original Images_rotated"
    rotate_and_save(image_dir, target_dir, angle_list, islabel = False)