import os
import argparse
from tqdm import tqdm
from pdf2image import convert_from_path

def main():
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='将PDF文件转换为图片')
    parser.add_argument('--input_dir', type=str, required=True, help='PDF文件输入目录')
    parser.add_argument('--output_dir', type=str, required=True, help='图片输出目录')
    
    args = parser.parse_args()
    
    # 确保输出目录存在
    os.makedirs(args.output_dir, exist_ok=True)
    
    # 获取PDF文件列表
    pdf_files = [file for file in os.listdir(args.input_dir) if file.endswith('.pdf')]
    
    for filename in tqdm(pdf_files):
        filepath = os.path.join(args.input_dir, filename)
        imgname = filename.split('.pdf')[0]
        images = convert_from_path(filepath)
        for i, image in enumerate(images):
            idx = i + 1
            image.save(os.path.join(args.output_dir, f'{imgname}_{idx}.jpg'), 'JPEG')

if __name__ == '__main__':
    main()