import glob
import os
import argparse
from PIL import Image


def make(img_dir, save_dir):
    output_path1 = 'img_270'
    output_path2 = 'img_180'
    output_path3 = 'img_90'
    output_path4 = 'img_0'
    if not os.path.exists(save_dir + '/' + output_path4):
        os.makedirs(save_dir + '/' + output_path4)
    if not os.path.exists(save_dir + '/' + output_path3):
        os.makedirs(save_dir + '/' + output_path3)
    if not os.path.exists(save_dir + '/' + output_path2):
        os.makedirs(save_dir + '/' + output_path2)
    if not os.path.exists(save_dir + '/' + output_path1):
        os.makedirs(save_dir + '/' + output_path1)
    img_list = os.listdir(img_dir)
    print(img_list)
    angle_list = []
    for i in range(0, 24):
        angle_list.append(15 * i)

    for img_path in img_list:
        print(img_dir + '/' + img_path)
        img_name = img_path.split('.')[0]
        im = Image.open(img_dir + '/' + img_path)
        im = im.convert("RGB")  # 把PNG格式转换成的四通道转成RGB的三通道
        fp = open('all_list.txt', 'a')
        for angle in angle_list:
            im_rotate = im.rotate(angle, expand=0, fillcolor='#FFFFFF')
            if 45 < angle <= 135:
                output_path = output_path1
                label = 3
            elif 135 < angle <= 225:
                output_path = output_path2
                label = 2
            elif 225 < angle <= 315:
                output_path = output_path3
                label = 1
            else:
                output_path = output_path4
                label = 0
            im_rotate.save(save_dir + '/' + output_path + '/' + img_name + '_rotated_' + str(angle) + '.jpg')
            fp.writelines(output_path + '/' + img_name + '_rotated_' + str(angle) + '.jpg ' + str(label) + '\n')
        fp.close()
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--img_dir', type=str, default='./img')
    parser.add_argument('--save_dir', type=str, default='./img_rotated')
    args = parser.parse_args()
    if not os.path.exists(args.save_dir):
        os.makedirs(args.save_dir)
    make(args.img_dir, args.save_dir)