225 lines
9.6 KiB
Python
225 lines
9.6 KiB
Python
import os
|
|
import sys
|
|
import json
|
|
import argparse
|
|
import logging
|
|
from typing import List, Dict
|
|
|
|
# 添加项目根目录到系统路径
|
|
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
|
|
# 导入所需的模块
|
|
from utils.output_handler import FileSystemOutputHandler
|
|
from utils.poster_notes_creator import select_additional_images
|
|
from utils.tweet_generator import generate_posters_for_topic
|
|
from core.topic_parser import TopicParser
|
|
|
|
def load_config(config_path="poster_gen_config.json"):
|
|
"""加载配置文件"""
|
|
if not os.path.exists(config_path):
|
|
logging.error(f"错误：配置文件 '{config_path}' 不存在")
|
|
return None
|
|
|
|
try:
|
|
with open(config_path, 'r', encoding='utf-8') as f:
|
|
config = json.load(f)
|
|
return config
|
|
except Exception as e:
|
|
logging.error(f"加载配置文件时出错: {e}")
|
|
return None
|
|
|
|
def regenerate_posters_and_images(run_dir: str, config: Dict):
|
|
"""为所有主题重新生成海报和配图"""
|
|
run_id = os.path.basename(run_dir)
|
|
logging.info(f"处理运行ID: {run_id}")
|
|
|
|
# 创建输出处理器
|
|
output_handler = FileSystemOutputHandler(config.get("output_dir", "result"))
|
|
|
|
# 加载主题文件
|
|
topics_file = os.path.join(run_dir, f"tweet_topic_{run_id}.json")
|
|
if not os.path.exists(topics_file):
|
|
alternative_topics_file = os.path.join(run_dir, "tweet_topic.json")
|
|
if not os.path.exists(alternative_topics_file):
|
|
logging.error(f"主题文件不存在: {topics_file} 或 {alternative_topics_file}")
|
|
return
|
|
topics_file = alternative_topics_file
|
|
|
|
logging.info(f"使用主题文件: {topics_file}")
|
|
topics_list = TopicParser.load_topics_from_json(topics_file)
|
|
if not topics_list:
|
|
logging.error("无法加载主题列表")
|
|
return
|
|
|
|
# 加载配置参数
|
|
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")
|
|
poster_content_system_prompt_path = config.get("poster_content_system_prompt")
|
|
collage_style = config.get("collage_style")
|
|
title_possibility = config.get("title_possibility", 0.3)
|
|
request_timeout = config.get("request_timeout", 180)
|
|
additional_images_count = config.get("additional_images_count", 3)
|
|
|
|
# 获取图像处理参数
|
|
image_selection_config = config.get("image_selection", {})
|
|
variation_strength = image_selection_config.get("variation_strength", "medium")
|
|
extra_effects = image_selection_config.get("extra_effects", True)
|
|
|
|
# 检查关键配置
|
|
if not poster_assets_dir or not img_base_dir or not poster_content_system_prompt_path:
|
|
logging.error("缺少关键配置参数")
|
|
return
|
|
|
|
# 读取海报内容系统提示词
|
|
with open(poster_content_system_prompt_path, "r", encoding="utf-8") as f:
|
|
poster_content_system_prompt = f.read()
|
|
|
|
# 处理每个主题
|
|
for topic_item in topics_list:
|
|
topic_index = topic_item.get('index')
|
|
if not topic_index:
|
|
logging.warning(f"主题缺少索引，跳过: {topic_item}")
|
|
continue
|
|
|
|
logging.info(f"处理主题 {topic_index}: {topic_item.get('object', 'N/A')}")
|
|
|
|
# 1. 为此主题生成所有变体的海报（一次性调用）
|
|
logging.info(f"为主题 {topic_index} 生成所有变体的海报...")
|
|
|
|
try:
|
|
posters_generated = 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,
|
|
model_name=config["model"],
|
|
base_url=config["api_url"],
|
|
api_key=config["api_key"],
|
|
variants=poster_variants, # 生成所有变体
|
|
title_possibility=title_possibility,
|
|
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,
|
|
timeout=request_timeout
|
|
)
|
|
|
|
if posters_generated:
|
|
logging.info(f"主题 {topic_index} 的所有海报生成完成")
|
|
else:
|
|
logging.warning(f"主题 {topic_index} 的海报生成失败或未返回成功标志")
|
|
continue
|
|
|
|
except Exception as e:
|
|
logging.exception(f"生成主题 {topic_index} 海报时出错: {e}")
|
|
continue
|
|
|
|
# 2. 为每个变体生成额外配图
|
|
for variant_index in range(1, poster_variants + 1):
|
|
variant_key = f"{topic_index}_{variant_index}"
|
|
topic_dir = os.path.join(run_dir, variant_key)
|
|
|
|
if not os.path.exists(topic_dir):
|
|
logging.warning(f"变体目录不存在: {topic_dir}，跳过生成配图")
|
|
continue
|
|
|
|
# 生成额外配图
|
|
logging.info(f"为变体 {variant_key} 生成额外配图...")
|
|
|
|
# 获取海报元数据
|
|
poster_dir = os.path.join(topic_dir, poster_subdir)
|
|
|
|
if not os.path.exists(poster_dir):
|
|
logging.warning(f"变体 {variant_key} 海报目录不存在，跳过生成配图")
|
|
continue
|
|
|
|
try:
|
|
# 查找元数据文件
|
|
metadata_files = [f for f in os.listdir(poster_dir) if f.endswith("_metadata.json")]
|
|
|
|
if not metadata_files:
|
|
logging.warning(f"变体 {variant_key} 未找到海报元数据文件，跳过生成配图")
|
|
continue
|
|
|
|
poster_metadata_path = os.path.join(poster_dir, metadata_files[0])
|
|
|
|
# 获取源图像目录
|
|
object_name = topic_item.get("object", "").split(".")[0].replace("景点信息-", "").strip()
|
|
source_image_dir = os.path.join(img_base_dir, object_name)
|
|
|
|
if not os.path.exists(source_image_dir) or not os.path.isdir(source_image_dir):
|
|
logging.warning(f"变体 {variant_key} 源图像目录不存在: {source_image_dir}")
|
|
continue
|
|
|
|
# 生成额外配图
|
|
additional_paths = select_additional_images(
|
|
run_id=run_id,
|
|
topic_index=topic_index,
|
|
variant_index=variant_index,
|
|
poster_metadata_path=poster_metadata_path,
|
|
source_image_dir=source_image_dir,
|
|
num_additional_images=additional_images_count,
|
|
output_handler=output_handler,
|
|
variation_strength=variation_strength,
|
|
extra_effects=extra_effects
|
|
)
|
|
|
|
if additional_paths:
|
|
logging.info(f"已为变体 {variant_key} 生成 {len(additional_paths)} 张额外配图")
|
|
else:
|
|
logging.warning(f"未能为变体 {variant_key} 生成任何额外配图")
|
|
|
|
except Exception as e:
|
|
logging.exception(f"生成变体 {variant_key} 额外配图时出错: {e}")
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="为所有主题重新生成海报和配图")
|
|
parser.add_argument("--run_dir", required=True, help="之前运行结果的目录")
|
|
parser.add_argument("--config", default="poster_gen_config.json", help="配置文件路径")
|
|
parser.add_argument("--debug", action="store_true", help="启用调试日志")
|
|
parser.add_argument("--topic", type=int, help="只处理指定主题索引")
|
|
args = parser.parse_args()
|
|
|
|
# 设置日志级别
|
|
log_level = logging.DEBUG if args.debug else logging.INFO
|
|
logging.basicConfig(
|
|
level=log_level,
|
|
format='%(asctime)s - %(levelname)s - [%(filename)s:%(lineno)d] - %(message)s',
|
|
datefmt='%Y-%m-%d %H:%M:%S'
|
|
)
|
|
|
|
# 加载配置
|
|
config = load_config(args.config)
|
|
if not config:
|
|
sys.exit(1)
|
|
|
|
# 检查目录是否存在
|
|
if not os.path.exists(args.run_dir):
|
|
logging.error(f"指定的运行目录不存在: {args.run_dir}")
|
|
sys.exit(1)
|
|
|
|
# 开始处理
|
|
logging.info(f"开始为 {args.run_dir} 重新生成海报和配图")
|
|
regenerate_posters_and_images(args.run_dir, config)
|
|
logging.info("处理完成")
|
|
|
|
if __name__ == "__main__":
|
|
main() |