# 提前处理图像加速训练速度
import os

from PIL import Image
from pathlib import Path
from tqdm import tqdm


def resize_image(data_root: str,
                 height: int = 512,
                 width: int = 384):
    data_root = Path(data_root)
    
    data_path_list = [*data_root.glob("*.png"), *data_root.glob("*.jpg"), *data_root.glob("*.jpeg"), *data_root.glob("*.JPEG")]
    
    new_data_root = str(data_root) + f"_{height}"
    
    os.makedirs(new_data_root, exist_ok=True)
    
    for data_path in tqdm(data_path_list):
        image_name = data_path.name
        save_path = os.path.join(new_data_root, image_name)
        image = Image.open(str(data_path))
        image = image.resize((width, height), Image.LANCZOS)
        image.save(save_path)


def main(args):
    resize_image(args.data_root, args.height, args.width)


if __name__ == "__main__":
    from argparse import ArgumentParser
    
    parser = ArgumentParser()
    
    parser.add_argument("--data_root", type=str)
    
    parser.add_argument("--height", type=int, default=512)
    
    parser.add_argument("--width", type=int, default=384)
    
    args = parser.parse_args()
    
    main(args)
    
    
# find . -type f -name "*_rendered.png" -exec bash -c 'mv "$0" "${0/_rendered.png/.png}"' {} \;