import cv2
import numpy as np
import albumentations as A
import os

# 入力フォルダ
input_path = "./tmp2"
# 出力フォルダ
output_path = "./tmp3"

# 拡張枚数の設定
generate_img = 5

img_size = 200

transform = A.Compose([
    # 反転（ランダム、確率50%）
    A.HorizontalFlip(p=0.5),
    # 回転（ランダム、確率100%）
    A.Rotate(limit=(-50, 50), p=1),
    # ランダム画像サイズ変更（スケーリング係数0.8～1.4、確率50%）
    A.RandomScale(scale_limit=(0.8, 1.2), p=0.8),
    # ランダムトリミングとリサイズ（リサイズ後のサイズ、トリミング領域、確率100%）
    ###A.RandomResizedCrop(height=img_size, width=img_size, scale=(1.0, 1.0), p=1)
    # albumentationsの新バージョンに対応
    A.RandomResizedCrop(size=(img_size, img_size), scale=(1.0, 1.0), p=1)
])

# 入力フォルダからファイル名を取得
all_files = os.listdir(input_path)
#特定の拡張子のファイルのみを取り出す
files = [i for i in all_files if i.endswith('.png') == True]

for file_name in files:

    # 画像読み込み
    img_path = input_path + '/' + file_name

    img = cv2.imread(img_path, cv2.IMREAD_COLOR)

    # ファイル名抽出
    file_name_1 = file_name.split('.')[0]

    for i in range(generate_img):
        # データ拡張
        cropped_image = transform(image=img)['image']

        # 画像保存
        cv2.imwrite(output_path + '/' + file_name_1 + '_' + str(i) + '.jpg', cropped_image)