#  图片压缩脚本
from PIL import Image
import os
from tqdm import tqdm
from io import BytesIO

def auto_compress_images(input_folder, output_folder, target_size_kb):
    """
    自动选择压缩质量以达到目标文件大小的函数
    
    :param input_folder: 输入图片文件夹路径
    :param output_folder: 输出图片文件夹路径
    :param target_size_kb: 目标文件大小 (KB)
    """
    # 创建输出文件夹
    os.makedirs(output_folder, exist_ok=True)
    
    # 遍历输入文件夹及其子文件夹中的所有图像文件
    image_files = []
    for root, _, files in os.walk(input_folder):
        for file in files:
            if file.lower().endswith(('.jpg', '.jpeg', '.png')):
                image_files.append((root, file))
    
    # 使用tqdm显示进度条
    for root, file in tqdm(image_files, desc="压缩进度"):
        input_path = os.path.join(root, file)
        output_path = os.path.join(output_folder, file)
        
        try:
            # 打开图像文件
            img = Image.open(input_path)

            # 检查并转换图像模式
            if img.mode not in ('RGB', 'L'):
                img = img.convert('RGB')

            # 初始质量
            quality = 80
            
            while True:
                # 在内存中保存压缩后的图像
                buffer = BytesIO()
                img.save(buffer, format='JPEG', quality=quality, optimize=True)
                
                # 获取压缩后的文件大小
                buffer_size_kb = len(buffer.getvalue()) / 1024
                
                if buffer_size_kb <= target_size_kb or quality <= 10:
                    with open(output_path, 'wb') as f:
                        f.write(buffer.getvalue())
                    break
                
                # 减小质量进一步压缩
                quality -= 5
                
        except Exception as e:
            print(f"处理图像 {file} 时出错: {e}")

# 示例用法
input_folder = './book/通用书籍/2024.8.16/'
output_folder = './book/通用书籍/2024.8.16/newImgs/'
target_size_kb = 100  # 目标文件大小 (KB)

auto_compress_images(input_folder, output_folder, target_size_kb)