PictureEdit/exporter.py

"""
作者：太一
微信：taiyi1224
邮箱：shuobo1224@qq.com
"""

from pathlib import Path
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor, as_completed
from PIL import Image
from image_manager import ImageManager
import os


class Exporter:
    """导出器，处理批量导出功能"""

    @staticmethod
    def export_batch(project, output_dir=None, filename_template="{name}_{idx}",
                     format='png', quality=95, keep_exif=False, progress_callback=None):
        """批量导出所有前景图"""
        # 如果没有指定输出目录，默认在第一个图片所在目录下创建output文件夹
        if output_dir is None and project.foregrounds:
            first_image_dir = Path(project.foregrounds[0].file_path).parent
            output_dir = first_image_dir / "output"
        
        output_dir = Path(output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)

        tasks = []
        for idx, fg in enumerate(project.foregrounds):
            base_name = Path(fg.file_path).stem
            filename = filename_template.format(
                name=base_name,
                idx=idx + 1,
                total=len(project.foregrounds),
                date=datetime.now().strftime("%Y%m%d")
            )
            file_path = output_dir / f"{filename}.{format.lower()}"
            tasks.append((fg, file_path))

        # 使用线程池执行导出任务，限制最大线程数为4，避免系统卡顿
        max_workers = min(4, os.cpu_count() or 1)  # 最多使用4个线程，或者CPU核心数（取较小值）
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            futures = []
            for fg, path in tasks:
                future = executor.submit(
                    Exporter._export_single,
                    project.template, fg, path, format, quality, keep_exif
                )
                future.add_done_callback(lambda f, i=len(futures):
                                         progress_callback(i + 1, len(tasks)) if progress_callback else None)
                futures.append(future)

            # 检查是否有错误
            errors = []
            for future in as_completed(futures):
                try:
                    result = future.result(timeout=30)  # 设置30秒超时
                    if not result:
                        errors.append(1)
                except Exception as e:
                    print(f"导出错误: {str(e)}")
                    errors.append(1)

        return len(errors) == 0

    @staticmethod
    def _export_single(template, foreground, output_path, format, quality, keep_exif):
        """导出单个前景图"""
        try:
            # 加载原图
            img = ImageManager.load_image(foreground.file_path)
            if img is None:
                return False

            # 应用遮罩
            masked_img = ImageManager.apply_mask(img, foreground.mask_shape)

            # 使用基于顶点的图像变形方法
            transformed_img, (x, y) = ImageManager.transform_with_vertices(
                masked_img, foreground.vertices
            )

            # 创建模板大小的画布
            canvas = Image.new('RGBA', (template.width_px, template.height_px),
                               color=template.bg_color + "FF")  # 添加不透明度

            # 绘制背景图
            if template.bg_image:
                # 缩放背景图以适应模板
                bg_scaled = template.bg_image.resize(
                    (template.width_px, template.height_px),
                    Image.Resampling.LANCZOS
                )
                canvas.paste(bg_scaled, (0, 0))

            # 粘贴前景图
            # 确保transformed_img具有RGBA模式
            if transformed_img.mode != 'RGBA':
                transformed_img = transformed_img.convert('RGBA')
            
            canvas.paste(transformed_img, (int(x), int(y)), transformed_img)

            # 保存图像
            return ImageManager.save_image(canvas, str(output_path), format, quality, keep_exif)

        except Exception as e:
            print(f"导出单个图像失败: {e}")
            return False