173 lines
5.9 KiB
Python
173 lines
5.9 KiB
Python
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("脚本执行完毕。") |