TravelContentCreator/scripts/generate_vibrant_poster.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
使用LLM和Vibrant模板从源文件自动生成海报的脚本。
"""

import os
import sys
import json
import json_repair
import random
import asyncio
import argparse
import logging
import re
from datetime import datetime
# 将项目根目录添加到Python路径中
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from poster.templates.vibrant_template import VibrantTemplate
from core.ai.ai_agent import AIAgent
from core.config import AIModelConfig
from multiprocessing.pool import ThreadPool
# 配置日志记录
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def read_data_file(file_path: str):
    """安全地读取并解析各种格式的文件。"""
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()
            
            # 尝试解析为标准JSON
            try:
                return json_repair.loads(content)
            except json.JSONDecodeError:
                # 如果不是标准JSON，尝试解析为类JSON格式
                logger.info(f"文件 {file_path} 不是标准JSON格式，尝试其他解析方式...")
                
                # 检查是否是类JSON格式（带有大括号的键值对）
                if content.strip().startswith('{') and content.strip().endswith('}'):
                    try:
                        # 尝试将内容解析为字典
                        result = {}
                        # 使用正则表达式提取键值对
                        pattern = r'"([^"]+)":\s*(.+?)(?=,\s*"[^"]+":|$)'
                        matches = re.findall(pattern, content, re.DOTALL)
                        
                        for key, value in matches:
                            # 清理值中的引号和多余空格
                            value = value.strip()
                            if value.startswith('"') and value.endswith('"'):
                                # 字符串值
                                result[key] = value[1:-1]
                            elif value.lower() in ['true', 'false']:
                                # 布尔值
                                result[key] = value.lower() == 'true'
                            elif value.replace('.', '', 1).isdigit():
                                # 数字值
                                result[key] = float(value) if '.' in value else int(value)
                            else:
                                # 其他值作为字符串处理
                                result[key] = value
                                
                        if result:
                            logger.info(f"成功解析文件 {file_path} 为字典")
                            return result
                    except Exception as e:
                        logger.error(f"尝试解析类JSON格式失败: {e}")
                
                # 如果上述方法失败，尝试直接提取文本内容
                logger.info(f"尝试将文件 {file_path} 作为纯文本处理")
                
                # 简单解析键值对格式的文本文件
                result = {}
                lines = content.split('\n')
                current_key = None
                current_value = []
                
                for line in lines:
                    line = line.strip()
                    if not line:
                        continue
                        
                    # 检查是否是新的键值对
                    if ':' in line and not line.startswith(' ') and not line.startswith('\t'):
                        # 保存之前的键值对
                        if current_key and current_value:
                            result[current_key] = '\n'.join(current_value)
                            current_value = []
                            
                        # 提取新的键值对
                        parts = line.split(':', 1)
                        current_key = parts[0].strip().strip('"').strip("'")
                        if len(parts) > 1 and parts[1].strip():
                            current_value = [parts[1].strip()]
                        else:
                            current_value = []
                    elif current_key:
                        # 继续添加到当前值
                        current_value.append(line)
                
                # 保存最后一个键值对
                if current_key and current_value:
                    result[current_key] = '\n'.join(current_value)
                
                if result:
                    logger.info(f"成功将文件 {file_path} 解析为键值对")
                    return result
                
                # 如果所有方法都失败，返回文本内容
                logger.warning(f"无法解析文件 {file_path} 为结构化数据，返回原始文本")
                return {"content": content}
                
    except FileNotFoundError:
        logger.error(f"文件未找到: {file_path}")
    except Exception as e:
        logger.error(f"读取文件时发生未知错误 {file_path}: {e}")
    return None

async def generate_info_with_llm(ai_agent: AIAgent, scenic_info: dict, product_info: dict, tweet_info: dict):
    """
    调用LLM，根据输入信息生成海报所需的vibrant_content。
    
    Args:
        ai_agent: AIAgent实例
        scenic_info: 景区信息字典
        product_info: 产品信息字典
        tweet_info: 推文信息字典
        
    Returns:
        一个符合VibrantTemplate格式的内容字典
    """
    logger.info("正在调用LLM以生成海报内容...")
    
    # 构建系统提示
    system_prompt = """你是一名专业的海报设计师，专门设计宣传海报。你现在要根据用户提供的图片描述和推文内容，生成适合模板的海报信息。

模板特点：
- 单图背景，毛玻璃渐变效果
- 两栏布局（左栏内容，右栏价格）
- 适合展示套餐内容和价格信息

你需要生成的数据结构包含以下字段：

**必填字段：**
1. `title`: 主标题（8-12字符，体现产品特色）
2. `slogan`: 副标题/宣传语（10-20字符，吸引人的描述）
3. `price`: 价格数字（纯数字，不含符号）
4. `ticket_type`: 票种类型（如"成人票"、"套餐票"、"夜场票"等）
5. `content_button`: 内容按钮文字（通常为"套餐内容"、"包含项目"等）
6. `content_items`: 套餐内容列表（3-5个项目，每项5-15字符，不要只包含项目名称，要做合适的美化，可以适当省略）

**可选字段：**
7. `remarks`: 备注信息（1-3条，每条10-20字符）
8. `tag`: 标签（1条, 如"#限时优惠"、""等）
9. `pagination`: 分页信息（如"1/3"，可为空）

**内容创作要求：**
1. 套餐内容要具体实用：明确说明包含的服务、时间、数量
2. 价格要有吸引力：突出性价比和优惠信息

**输出格式：**
请严格按照以下JSON格式输出，不要有任何额外内容：

```json
{
    "title": "海洋奇幻世界",
    "slogan": "探索深海秘境，感受蓝色奇迹的无限魅力",
    "price": "299",
    "ticket_type": "成人票",
    "content_button": "套餐内容",
    "content_items": [
        "海洋馆门票1张（含所有展区）",
        "海豚表演VIP座位",
        "鲨鱼隧道特别体验",
        "专业摄影服务"
    ],
    "remarks": [
        "工作日可直接入园",
        "周末需提前预约"
    ],
    "tag": "#海洋特惠",
    "pagination": "1/2"
}
```
"""

    # 构建用户提示
    user_prompt = f"""请根据以下信息，生成适合在旅游海报上展示的文案：

## 景区信息
{json.dumps(scenic_info, ensure_ascii=False, indent=2)}

## 产品信息
{json.dumps(product_info, ensure_ascii=False, indent=2)}

## 推文信息
{json.dumps(tweet_info, ensure_ascii=False, indent=2)}

请提取关键信息并整合成一个JSON对象，包含title、slogan、price、ticket_type、content_items、remarks和tag字段。
"""
    
    try:
        # 调用AIAgent生成文本
        response, _, _, _ = await ai_agent.generate_text(
            system_prompt=system_prompt,
            user_prompt=user_prompt,
            temperature=0.7,
            stage="海报文案生成"
        )
        
        # 解析JSON响应
        try:
            # 尝试找到JSON对象的开始和结束位置
            json_start = response.find('{')
            json_end = response.rfind('}') + 1
            
            if json_start >= 0 and json_end > json_start:
                json_str = response[json_start:json_end]
                content_dict = json_repair.loads(json_str)
                logger.info(f"LLM成功生成内容: {content_dict}")
                
                # 添加默认的按钮文本和分页信息
                content_dict["content_button"] = content_dict.get("content_button", "套餐内容")
                content_dict["pagination"] = content_dict.get("pagination", "")
                
                # 确保所有值都是字符串类型
                for key, value in content_dict.items():
                    if isinstance(value, (int, float)):
                        content_dict[key] = str(value)
                    elif isinstance(value, list):
                        content_dict[key] = [str(item) if isinstance(item, (int, float)) else item for item in value]
                
                logger.info(f"转换类型后的内容: {content_dict}")
                return content_dict
            else:
                logger.error(f"无法在响应中找到JSON对象: {response}")
                return None
        except json.JSONDecodeError:
            logger.error(f"无法解析LLM响应为JSON: {response}")
            return None
            
    except Exception as e:
        logger.error(f"调用LLM时发生错误: {e}")
        return None

async def generate_poster(img_dir, output_dir, scenic_info_file, product_info_file, tweet_info_file, model, api_key, api_url):
    """
    根据提供的参数生成海报。
    """
    # --- 1. 读取所有源文件 ---
    scenic_info = read_data_file(scenic_info_file)
    product_info = read_data_file(product_info_file)
    tweet_info = read_data_file(tweet_info_file)

    if not all([scenic_info, product_info, tweet_info]):
        logger.error("一个或多个源文件读取失败，终止执行。")
        return

    # --- 2. 初始化AI代理 ---
    ai_config = AIModelConfig(
        model=model,
        api_key=api_key,
        api_url=api_url
    )
    ai_agent = AIAgent(ai_config)

    # --- 3. 调用LLM生成海报内容 ---
    vibrant_content = await generate_info_with_llm(ai_agent, scenic_info, product_info, tweet_info)
    if not vibrant_content:
        logger.error("未能从LLM获取有效内容，终止执行。")
        return

    # --- 4. 随机选择图片 ---
    try:
        image_files = [f for f in os.listdir(img_dir) if f.lower().endswith(('png', 'jpg', 'jpeg', 'webp'))]
        if not image_files:
            logger.error(f"在目录 {img_dir} 中未找到任何图片文件。")
            return
        random_image_name = random.choice(image_files)
        image_path = os.path.join(img_dir, random_image_name)
        logger.info(f"随机选择图片: {image_path}")
    except FileNotFoundError:
        logger.error(f"图片目录不存在: {img_dir}")
        return
        
    # --- 5. 调用模板生成海报 ---
    try:
        logger.info("正在初始化 VibrantTemplate...")
        template = VibrantTemplate()
        poster = template.generate(image_path=image_path, content=vibrant_content)
        
        if not poster:
            logger.error("海报生成失败，模板返回了 None。")
            return
            
    except Exception as e:
        logger.error(f"生成海报时发生未知错误: {e}", exc_info=True)
        return

    # --- 6. 保存图片 ---
    try:
        os.makedirs(output_dir, exist_ok=True)
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        scenic_name = scenic_info.get('主体名称', 'poster')
        output_filename = f"vibrant_{scenic_name}_{timestamp}.png"
        output_path = os.path.join(output_dir, output_filename)
        
        poster.save(output_path, 'PNG')
        logger.info(f"海报已成功生成并保存至: {output_path}")
        
    except Exception as e:
        logger.error(f"保存海报失败: {e}", exc_info=True)

def run_async(func, *args):  # 同步包装器
    return asyncio.run(func(*args))

async def main():
    parser = argparse.ArgumentParser(
        description="通过LLM从源文件生成信息，并使用Vibrant模板创建海报。",
        formatter_class=argparse.RawTextHelpFormatter
    )

    # 设置默认值
    img_dir = "/root/autodl-tmp/TCC_RESTRUCT/resource/data/images/天津冒险湾"
    output_dir = "result/天津冒险湾"
    scenic_info_file = "resource/data/Object/天津冒险湾.txt"
    product_info_file = "resource/data/Product/天津冒险湾-2大2小套票.txt"
    tweet_info_file = "/root/autodl-tmp/TCC_RESTRUCT/result/run_20250709_160942/topic_1/article_judged.json"
    
    model = "qwen-plus"
    api_key = "sk-bd5ee62703bc41fc9b8a55d748dc1eb8"
    api_url = "https://dashscope.aliyuncs.com/compatible-mode/v1"

    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    tasks = [
        (img_dir, output_dir, scenic_info_file, product_info_file, tweet_info_file, model, api_key, api_url),
        (img_dir, output_dir, scenic_info_file, product_info_file, tweet_info_file, model, api_key, api_url),
        (img_dir, output_dir, scenic_info_file, product_info_file, tweet_info_file, model, api_key, api_url),
        (img_dir, output_dir, scenic_info_file, product_info_file, tweet_info_file, model, api_key, api_url),
        (img_dir, output_dir, scenic_info_file, product_info_file, tweet_info_file, model, api_key, api_url),
    ]
    with ThreadPool(5) as pool:
        # 直接传递任务参数，不需要额外包装
        results = pool.starmap(
            run_async, 
            [(generate_poster, *task_args) for task_args in tasks]
        )
        pool.close()
        pool.join()

print("所有海报生成完成。")


if __name__ == "__main__":
    asyncio.run(main())