import cv2 as cv
import os
import tqdm


# 获取目录下所有文件的路径
def get_img_names(directory):
    file_paths = []
    for root, dirs, files in os.walk(directory):
        for file in files:
            file_paths.append(os.path.join(root, file))
    return file_paths


# 导入并预处理
def pretrain(img_path, save_path):
    img = cv.imread(img_path, cv.IMREAD_GRAYSCALE)
    _,img = cv.threshold(img,0,255,cv.THRESH_BINARY+cv.THRESH_OTSU)
    img = cv.blur(img, (3, 3))

    os.makedirs(os.path.dirname(save_path), exist_ok=True)

    cv.imwrite(save_path, img)

if __name__ == '__main__':
    trains_paths=get_img_names('DataImages-Train/')
    for train_path in tqdm.tqdm(trains_paths,desc='预处理训练文件中：'):
        save_path=os.path.join('cache/pretrains/train', os.path.basename(train_path))
        pretrain(train_path,save_path)

    tests_paths=get_img_names('DataImages-Test/')
    for test_path in tqdm.tqdm(tests_paths,desc='预处理测试文件中：'):
        save_path=os.path.join('cache/pretrains/test', os.path.basename(test_path))
        pretrain(test_path,save_path)