239 lines
9.2 KiB
Python
239 lines
9.2 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
测试海报生成的脚本
|
||
"""
|
||
|
||
import os
|
||
import time
|
||
import argparse
|
||
import json
|
||
import logging
|
||
import sys
|
||
import traceback
|
||
from datetime import datetime
|
||
|
||
from core.ai_agent import AI_Agent
|
||
from utils.tweet_generator import generate_posters_for_topic
|
||
from utils.output_handler import FileSystemOutputHandler
|
||
from core.topic_parser import TopicParser
|
||
|
||
def load_config(config_path="poster_config.json"):
|
||
"""从JSON文件加载配置"""
|
||
if not os.path.exists(config_path):
|
||
print(f"错误：配置文件 '{config_path}' 未找到。")
|
||
sys.exit(1)
|
||
try:
|
||
with open(config_path, 'r', encoding='utf-8') as f:
|
||
config = json.load(f)
|
||
|
||
# 基本验证
|
||
required_keys = ["api_url", "model", "api_key", "resource_dir", "output_dir",
|
||
"image_base_dir", "poster_assets_base_dir",
|
||
"poster_content_system_prompt"]
|
||
|
||
if not all(key in config for key in required_keys):
|
||
missing_keys = [key for key in required_keys if key not in config]
|
||
print(f"错误：配置文件 '{config_path}' 缺少必需的键：{missing_keys}")
|
||
sys.exit(1)
|
||
|
||
return config
|
||
except json.JSONDecodeError:
|
||
print(f"错误：无法从 '{config_path}' 解码JSON。请检查文件格式。")
|
||
sys.exit(1)
|
||
except Exception as e:
|
||
print(f"从 '{config_path}' 加载配置时出错：{e}")
|
||
sys.exit(1)
|
||
|
||
def main():
|
||
# 设置日志记录
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format='%(asctime)s - %(levelname)s - [%(filename)s:%(lineno)d] - %(message)s',
|
||
datefmt='%Y-%m-%d %H:%M:%S'
|
||
)
|
||
|
||
# 解析命令行参数
|
||
parser = argparse.ArgumentParser(description="测试海报生成")
|
||
parser.add_argument(
|
||
"--config",
|
||
type=str,
|
||
default="/root/autodl-tmp/TravelContentCreator/poster_gen_config.json",
|
||
help="配置文件路径（例如，poster_config.json）"
|
||
)
|
||
parser.add_argument(
|
||
"--topics_file",
|
||
type=str,
|
||
required=True,
|
||
default="/root/autodl-tmp/TravelContentCreator/result/2025-04-25_17-27-03/tweet_topic_2025-04-25_17-27-03.json",
|
||
help="必需的选题JSON文件路径，用于获取海报生成的主题数据。"
|
||
)
|
||
parser.add_argument(
|
||
"--topic_index",
|
||
type=int,
|
||
default=1,
|
||
help="要生成海报的特定选题索引。如果未提供，将为所有选题生成海报。"
|
||
)
|
||
parser.add_argument(
|
||
"--run_id",
|
||
type=str,
|
||
default=None,
|
||
help="可选的指定运行ID。如果未提供，将生成基于时间戳的ID。"
|
||
)
|
||
parser.add_argument(
|
||
"--debug",
|
||
action='store_true',
|
||
help="启用调试级别日志记录。"
|
||
)
|
||
args = parser.parse_args()
|
||
|
||
# 调整日志级别（如果启用了调试）
|
||
if args.debug:
|
||
logging.getLogger().setLevel(logging.DEBUG)
|
||
logging.info("已启用调试日志记录。")
|
||
|
||
logging.info("启动海报生成测试脚本...")
|
||
logging.info(f"使用配置文件：{args.config}")
|
||
logging.info(f"使用选题文件：{args.topics_file}")
|
||
if args.topic_index is not None:
|
||
logging.info(f"将仅处理选题索引：{args.topic_index}")
|
||
if args.run_id:
|
||
logging.info(f"使用指定的run_id：{args.run_id}")
|
||
|
||
# 加载配置
|
||
config = load_config(args.config)
|
||
if config is None:
|
||
logging.critical("无法加载配置。退出。")
|
||
sys.exit(1)
|
||
|
||
# 初始化输出处理器
|
||
output_handler = FileSystemOutputHandler(config.get("output_dir", "result"))
|
||
logging.info(f"使用输出处理器：{output_handler.__class__.__name__}")
|
||
|
||
# 加载选题数据
|
||
logging.info(f"从以下位置加载选题：{args.topics_file}")
|
||
topics_list = TopicParser.load_topics_from_json(args.topics_file)
|
||
if not topics_list:
|
||
logging.error(f"无法从{args.topics_file}加载选题。无法继续。")
|
||
sys.exit(1)
|
||
|
||
logging.info(f"成功加载{len(topics_list)}个选题。")
|
||
|
||
# 设置run_id
|
||
run_id = args.run_id
|
||
if run_id is None:
|
||
# 尝试从文件名推断run_id
|
||
try:
|
||
base = os.path.basename(args.topics_file)
|
||
if base.startswith("tweet_topic_") and base.endswith(".json"):
|
||
run_id = base[len("tweet_topic_"): -len(".json")]
|
||
logging.info(f"从选题文件名推断的run_id：{run_id}")
|
||
else:
|
||
run_id = datetime.now().strftime("%Y-%m-%d_%H-%M-%S_poster")
|
||
logging.info(f"为海报生成的run_id：{run_id}")
|
||
except Exception as e:
|
||
run_id = datetime.now().strftime("%Y-%m-%d_%H-%M-%S_poster")
|
||
logging.info(f"生成的run_id：{run_id}")
|
||
|
||
# 加载海报内容系统提示词
|
||
poster_content_system_prompt_path = config.get("poster_content_system_prompt")
|
||
if not os.path.exists(poster_content_system_prompt_path):
|
||
logging.error(f"海报内容系统提示词文件不存在：{poster_content_system_prompt_path}")
|
||
sys.exit(1)
|
||
|
||
with open(poster_content_system_prompt_path, "r", encoding="utf-8") as f:
|
||
poster_content_system_prompt = f.read()
|
||
|
||
# 准备海报生成参数
|
||
poster_variants = config.get("variants", 1)
|
||
poster_assets_dir = config.get("poster_assets_base_dir")
|
||
img_base_dir = config.get("image_base_dir")
|
||
res_dir_config = config.get("resource_dir", [])
|
||
poster_size = tuple(config.get("poster_target_size", [900, 1200]))
|
||
txt_possibility = config.get("text_possibility", 0.3)
|
||
img_frame_possibility = config.get("img_frame_possibility", 0.7)
|
||
text_bg_possibility = config.get("text_bg_possibility", 0)
|
||
collage_subdir = config.get("output_collage_subdir", "collage_img")
|
||
poster_subdir = config.get("output_poster_subdir", "poster")
|
||
poster_filename = config.get("output_poster_filename", "poster.jpg")
|
||
title_possibility = config.get("title_possibility", 0.3)
|
||
collage_style = config.get("collage_style", None)
|
||
timeout = config.get("request_timeout", 180)
|
||
|
||
# 检查关键路径
|
||
if not poster_assets_dir or not img_base_dir:
|
||
logging.error("配置中缺少关键路径（poster_assets_base_dir或image_base_dir）。无法继续。")
|
||
sys.exit(1)
|
||
|
||
# 开始海报生成
|
||
pipeline_start_time = time.time()
|
||
logging.info("开始执行海报生成...")
|
||
|
||
poster_success = False
|
||
|
||
# 如果指定了topic_index，只处理该选题
|
||
if args.topic_index is not None:
|
||
topics_to_process = []
|
||
for topic in topics_list:
|
||
if topic.get('index') == args.topic_index or (topic.get('index') is None and int(args.topic_index) == 1):
|
||
topics_to_process.append(topic)
|
||
break
|
||
if not topics_to_process:
|
||
logging.error(f"未找到索引为{args.topic_index}的选题。")
|
||
sys.exit(1)
|
||
else:
|
||
topics_to_process = topics_list
|
||
|
||
# 逐个处理选题
|
||
for i, topic_item in enumerate(topics_to_process):
|
||
topic_index = topic_item.get('index', i + 1)
|
||
logging.info(f"--- 处理选题 {topic_index}: {topic_item.get('object', 'N/A')} ---")
|
||
|
||
try:
|
||
posters_attempted = generate_posters_for_topic(
|
||
topic_item=topic_item,
|
||
output_dir=config["output_dir"],
|
||
run_id=run_id,
|
||
topic_index=topic_index,
|
||
output_handler=output_handler,
|
||
variants=poster_variants,
|
||
poster_assets_base_dir=poster_assets_dir,
|
||
image_base_dir=img_base_dir,
|
||
resource_dir_config=res_dir_config,
|
||
poster_target_size=poster_size,
|
||
text_possibility=txt_possibility,
|
||
img_frame_possibility=img_frame_possibility,
|
||
text_bg_possibility=text_bg_possibility,
|
||
output_collage_subdir=collage_subdir,
|
||
output_poster_subdir=poster_subdir,
|
||
output_poster_filename=poster_filename,
|
||
system_prompt=poster_content_system_prompt,
|
||
collage_style=collage_style,
|
||
title_possibility=title_possibility,
|
||
timeout=timeout
|
||
)
|
||
|
||
if posters_attempted:
|
||
logging.info(f"选题{topic_index}的海报生成过程已完成。")
|
||
poster_success = True
|
||
else:
|
||
logging.warning(f"选题{topic_index}的海报生成被跳过或在早期失败。")
|
||
except Exception as e:
|
||
logging.exception(f"处理选题{topic_index}的海报生成时出错:")
|
||
logging.error(f"错误信息: {e}")
|
||
logging.info(f"--- 完成选题 {topic_index} ---")
|
||
|
||
# 最终化输出
|
||
if run_id:
|
||
output_handler.finalize(run_id)
|
||
|
||
pipeline_end_time = time.time()
|
||
if poster_success:
|
||
logging.info(f"海报生成完成，耗时{pipeline_end_time - pipeline_start_time:.2f}秒。")
|
||
else:
|
||
logging.warning("海报生成完成，但可能遇到错误或未生成输出。")
|
||
|
||
logging.info(f"运行ID'{run_id}'的结果位于：{os.path.join(config.get('output_dir', 'result'), run_id)}")
|
||
|
||
if __name__ == "__main__":
|
||
main() |