TravelContentCreator/scripts/select_additional_images.py

import argparse
import logging
from pathlib import Path
import sys
import os # Needed for path checks

# --- Path Setup ---
# Add project root to sys.path to allow importing utils
project_root = Path(__file__).resolve().parent.parent
if str(project_root) not in sys.path:
    sys.path.insert(0, str(project_root))
    print(f"Added project root {project_root} to sys.path")

try:
    # Import the necessary function and class from your utils
    from utils.poster_notes_creator import select_additional_images
    from utils.output_handler import OutputHandler
except ImportError as e:
    print(f"错误: 无法导入所需的模块 'utils.poster_notes_creator' 或 'utils.output_handler'.")
    print(f"错误详情: {e}")
    print(f"Python 搜索路径 (sys.path): {sys.path}")
    print("请确保你在项目根目录下运行此脚本，或者 utils 目录在 PYTHONPATH 中。")
    sys.exit(1)

try:
    # Check for Pillow dependency implicitly used by poster_notes_creator
    from PIL import Image
except ImportError:
    print("错误: 未找到所需的库 PIL (Pillow)。")
    print("请先安装它: pip install Pillow")
    sys.exit(1)


# --- Logging Configuration ---
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - [%(filename)s:%(lineno)d] - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)

# --- 主程序入口 ---

if __name__ == "__main__":
    # 设置命令行参数解析
    parser = argparse.ArgumentParser(
        description="调用 poster_notes_creator 选择并处理额外配图。"
    )
    parser.add_argument(
        "--run_id",
        type=str,
        required=True,
        help="当前处理流程的运行 ID。"
    )
    parser.add_argument(
        "--topic_index",
        type=int,
        required=True,
        help="当前处理的主题索引。"
    )
    parser.add_argument(
        "--variant_index",
        type=int,
        required=True,
        help="当前处理的变体索引。"
    )
    parser.add_argument(
        "--metadata_path",
        type=str,
        required=True,
        help="指向包含已用图片列表的海报元数据 JSON 文件路径 (通常是 poster_metadata.json)。"
    )
    parser.add_argument(
        "--source_image_dir",
        type=str,
        required=True,
        help="包含源图片的目录路径。"
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        required=True,
        help="用于保存生成结果的根目录。"
    )
    parser.add_argument(
        "--num_additional",
        type=int,
        default=3,
        help="要选择的额外配图数量 (默认为 3)。"
    )
    parser.add_argument(
        "--variation_strength",
        type=str,
        default="medium",
        choices=["low", "medium", "high"],
        help="应用于额外配图的微调强度 (默认为 'medium')。"
    )
    parser.add_argument(
        "--extra_effects",
        action=argparse.BooleanOptionalAction, # 使用 --extra-effects 或 --no-extra-effects
        default=True,
        help="是否为额外配图添加额外视觉效果 (例如噪点、微调) (默认为启用)。"
    )
    parser.add_argument(
        "--debug",
        action='store_true',
        help="启用 DEBUG 级别日志。"
    )

    # 解析参数
    args = parser.parse_args()

    # 设置日志级别
    if args.debug:
        logging.getLogger().setLevel(logging.DEBUG)
        logger.info("DEBUG 日志已启用")

    # 输入验证
    metadata_file = Path(args.metadata_path)
    source_dir = Path(args.source_image_dir)
    output_root_dir = Path(args.output_dir)

    if not metadata_file.is_file():
        logger.error(f"指定的元数据文件不存在或不是一个文件: {metadata_file}")
        sys.exit(1)
    if not source_dir.is_dir():
        logger.error(f"指定的源图片目录不存在或不是一个目录: {source_dir}")
        sys.exit(1)
    # output_dir 的存在性由 OutputHandler 内部处理，这里不检查

    if args.num_additional <= 0:
        logger.error("选择的额外配图数量必须是一个正整数。")
        sys.exit(1)

    # --- 执行操作 ---
    try:
        # 1. 初始化 OutputHandler
        # 注意：output_handler 需要 output_dir 参数
        output_handler = OutputHandler(output_dir=str(output_root_dir))
        logger.info(f"OutputHandler 初始化完成，输出根目录: {output_root_dir}")

        # 2. 调用 poster_notes_creator 中的函数
        logger.info("开始调用 select_additional_images...")
        saved_image_paths = select_additional_images(
            run_id=args.run_id,
            topic_index=args.topic_index,
            variant_index=args.variant_index,
            poster_metadata_path=str(metadata_file),
            source_image_dir=str(source_dir),
            num_additional_images=args.num_additional,
            output_handler=output_handler, # 传递实例
            # output_filename_template 保持默认或根据需要添加参数
            variation_strength=args.variation_strength,
            extra_effects=args.extra_effects
        )

        # 3. 打印结果
        if saved_image_paths:
            print("\n========== 额外配图生成结果 ==========")
            for path in saved_image_paths:
                print(f"  已保存额外配图: {path}")
            print("===================================\n")
            logger.info(f"成功生成并保存了 {len(saved_image_paths)} 张额外配图。")
        else:
            print("\n未能生成或保存任何额外配图。请检查日志获取详细信息。")
            logger.warning("select_additional_images 函数没有返回任何有效的图片路径。")

    except Exception as e:
        logger.exception("执行额外配图选择和处理时发生严重错误:")
        print(f"脚本执行过程中发生错误，请查看日志。错误信息: {e}")
        sys.exit(1)

    logger.info("脚本执行完毕。")