import cv2 as cv
import os
import tqdm

# 处理图像
def pretrain(img_path,output_path):
    img=cv.imread(img_path, cv.IMREAD_GRAYSCALE)
    img = cv.resize(img, (50, 50))
    _,img=cv.threshold(img,127,255,cv.THRESH_BINARY)
    img=cv.blur(img,(3,3))
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    cv.imwrite(output_path, img)

# 获取文件路径
def get_img_name(directory, extensions=None):
    if extensions is None:
        extensions = ['.png', '.jpg','.jpeg']
    files = []
    for root, dirs, file_names in os.walk(directory):
        for file_name in file_names:
            if any(file_name.lower().endswith(ext) for ext in extensions):
                files.append(os.path.join(root, file_name))
    return files

# 处理非人脸数据集
os.makedirs('cache/pretrained/non_face/', exist_ok=True)
non_face_files=get_img_name('cache/dataset/non_face/')
print('预处理非人脸数据中：')
for img_path in tqdm.tqdm(non_face_files):
    pretrain(img_path, os.path.join('cache/pretrained/non_face', os.path.basename(img_path)))

# 处理人脸数据集
os.makedirs('cache/pretrained/face/', exist_ok=True)
face_files=get_img_name('cache/dataset/face/')
print('预处理人脸数据中：')
for img_path in tqdm.tqdm(face_files):
    relative_path=os.path.relpath(img_path, 'cache/dataset/face/')
    output_path=os.path.join('cache/pretrained/face', relative_path)
    pretrain(img_path, output_path)