TravelContentCreator/test_poster.py

#!/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()