from PIL import Image
import os
import shutil
from multiprocessing import Pool
from functools import partial

def convert_image(file_name, img_dir, output_dir, convert_format):
    input_path = os.path.join(img_dir, file_name)
    output_path = os.path.join(output_dir, file_name.replace(".jpg", f".{convert_format}"))
    
    # 打开并转换图像
    img = Image.open(input_path)
    img.save(output_path)

def main():
    convert_format = "png"
    img_dir = r"E:\datasets\UAV\134\project\images"
    output_dir = r"E:\datasets\UAV\134_png\project\images"

    # 如果输出目录存在，先删除
    if os.path.exists(output_dir):
        shutil.rmtree(output_dir)

    # 创建输出目录
    os.makedirs(output_dir)

    # 获取所有文件名
    file_names = os.listdir(img_dir)
    
    # 创建部分函数，固定除文件名外的其他参数
    convert_partial = partial(convert_image, 
                            img_dir=img_dir, 
                            output_dir=output_dir, 
                            convert_format=convert_format)

    # 使用进程池并行处理
    with Pool() as pool:
        pool.map(convert_partial, file_names)

if __name__ == '__main__':
    main()