import numpy as np
import os
import cv2

source_file = "D:/random_erasing_testdata4/"  # 原始文件地址
target_file = "D:/random_erasing_testdata4_output18/"  # 修改后的文件地址
num = 424  # 产生图片次数


def random_erasing(img, p=1, sl=0.01, sh=0.2, r1=0.3, r2=0.5):
    target_img = img.copy()

    if p < np.random.rand():
        # RandomErasingを実行しない
        return target_img

    H, W, _ = target_img.shape
    S = H * W

    while True:
        Se = np.random.uniform(sl, sh) * S  # 画像に重畳する矩形の面積
        re = np.random.uniform(r1, r2)  # 画像に重畳する矩形のアスペクト比

        He = int(np.sqrt(Se * re))  # 画像に重畳する矩形のHeight
        We = int(np.sqrt(Se / re))  # 画像に重畳する矩形のWidth

        # choose = np.random.randint(0, 2)
        # print(choose)
        # if choose == 0:
        #     xe = np.random.randint(0, W/3)  # 画像に重畳する矩形のx座標
        # if choose == 1:
        #     xe = np.random.randint(W*2/3, W)  # 画像に重畳する矩形のx座標

        xe = np.random.randint(0, W)  # 画像に重畳する矩形のx座標
        ye = np.random.randint(0, H/5)  # 画像に重畳する矩形のy座標

        if xe + We <= W and ye + He <= H:
            # 画像に重畳する矩形が画像からはみ出していなければbreak
            break

    mask = np.random.randint(0, 255, (He, We, 3)) # 矩形がを生成 矩形内の値はランダム値
    target_img[ye:ye + He, xe:xe + We, :] = mask # 画像に矩形を重畳

    return target_img


if not os.path.exists(target_file):  # 如果不存在target_file，则创造一个
    os.makedirs(target_file)


for i in range(num):
    print("Processing: ", str(i))
    file_list = os.listdir(source_file)
    img = cv2.imread(source_file + file_list[i])
    new_img = random_erasing(img)
    cv2.imwrite(target_file + file_list[i], new_img)