"""
FFHQ 1024 x 1024 -> 512 x 512
"""
import os

from zipfile import ZipFile
from PIL import Image
from glob import glob
from tqdm import tqdm


# 解压缩并转换图片大小
def process_in_zip(zip_dir: str,
                   output_dir: str):
    for zfile in glob(os.path.join(zip_dir, "*.zip")):
        with ZipFile(zfile, "r") as zip_ref:
            print("extract from ", zfile)
            for file_info in tqdm(zip_ref.infolist()):
                file_name = file_info.filename
                
                if file_name.lower().endswith(('.png', 'jpg', 'jpeg')):
                    with zip_ref.open(file_name) as file:
                        ori_image = Image.open(file)
                        
                        resized_image = ori_image.resize((512, 512), Image.LANCZOS)
                        
                        output_path = os.path.join(output_dir, file_name)
                        
                        os.makedirs(os.path.dirname(output_path), exist_ok=True)
                        
                        resized_image.save(output_path)


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    
    parser.add_argument("--zip_dir", type=str)
    
    parser.add_argument("--output_dir", type=str)
    
    args = parser.parse_args()
    
    process_in_zip(args.zip_dir, args.output_dir)