TravelContentCreator/scripts/test_poster_smart.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
测试智能海报生成引擎
"""

import asyncio
import base64
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

from domain.aigc.engines.poster_smart_v1 import PosterSmartEngine
from poster_v2 import PosterFactory, PosterContent

OUTPUT_DIR = Path(__file__).parent.parent / "result" / "poster_smart_test"
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

# 所有布局
ALL_LAYOUTS = ["hero_bottom", "overlay_center", "overlay_bottom", "split_vertical", "card_float"]
ALL_THEMES = ["ocean", "sunset", "peach", "mint", "latte"]


async def test_smart_poster():
    print("=" * 60)
    print("测试智能海报生成引擎")
    print("=" * 60)
    
    engine = PosterSmartEngine()
    
    # 测试用例
    test_cases = [
        {
            "name": "景点_正佳海洋世界",
            "params": {
                "category": "景点",
                "name": "正佳极地海洋世界",
                "description": "位于广州正佳广场的大型海洋馆，有企鹅、海豚表演，超适合带小朋友来",
                "price": "199元/人",
                "location": "广州天河",
                "features": "企鹅馆, 海豚表演, 儿童乐园, 室内恒温",
                "target_audience": "亲子家庭",
            }
        },
        {
            "name": "美食_网红甜品店",
            "params": {
                "category": "美食",
                "name": "Lady M 千层蛋糕",
                "description": "纽约网红甜品店，招牌千层蛋糕颜值和味道都在线",
                "price": "88元/份",
                "location": "深圳万象城",
                "features": "网红打卡, 颜值超高, 下午茶首选, 适合拍照",
                "target_audience": "闺蜜约会",
                "style_hint": "精致、少女感",
            }
        },
        {
            "name": "酒店_三亚度假酒店",
            "params": {
                "category": "酒店",
                "name": "三亚亚特兰蒂斯",
                "description": "超大水族馆酒店，推开窗就能看到鲨鱼游过",
                "price": "2888元/晚",
                "location": "三亚海棠湾",
                "features": "水族馆, 无边泳池, 私人沙滩, 水上乐园",
                "target_audience": "情侣度假",
            }
        },
        {
            "name": "民宿_山间咖啡民宿",
            "params": {
                "category": "民宿",
                "name": "山舍·云端民宿",
                "description": "藏在山里的宝藏民宿，坐在露台上看云海日出",
                "price": "458元/晚",
                "location": "莫干山",
                "features": "独立庭院, 手冲咖啡, 山景露台, 有机早餐",
                "target_audience": "周末度假",
                "style_hint": "自然、治愈",
            }
        },
        {
            "name": "活动_周末露营",
            "params": {
                "category": "活动",
                "name": "星空露营派对",
                "description": "周末一起去露营吧！篝火晚会、烧烤、看星星",
                "price": "299元/人",
                "location": "从化流溪河",
                "features": "篝火晚会, 星空观测, 烧烤BBQ, 帐篷已搭好",
                "target_audience": "年轻人",
            }
        },
        {
            "name": "攻略_广州三日游",
            "params": {
                "category": "攻略",
                "name": "广州三日游超全攻略",
                "description": "本地人带你玩转广州，景点美食一网打尽",
                "location": "广州",
                "features": "路线规划, 美食推荐, 避坑指南, 省钱技巧",
            }
        },
    ]
    
    # 测试带真实图片的
    test_image_path = Path("/root/TravelContentCreator/data/images/广东旅游_放弃伊犁这藏在广东的小新疆更有性价比_6836f1aa0000000021018c5a_1.jpg")
    
    print(f"\n[0] 测试图片: {test_image_path.name}")
    print(f"    存在: {test_image_path.exists()}")
    
    for i, case in enumerate(test_cases, 1):
        print(f"\n[{i}] 测试: {case['name']}")
        print(f"    类型: {case['params'].get('category')}")
        print(f"    名称: {case['params'].get('name')}")
        
        result = await engine.execute(case['params'])
        
        if result.success:
            layout = result.data.get('layout')
            theme = result.data.get('theme')
            content = result.data.get('generated_content', {})
            
            print(f"    ✓ 成功!")
            print(f"    布局: {layout}, 主题: {theme}")
            print(f"    标题: {content.get('title', 'N/A')}")
            print(f"    副标题: {content.get('subtitle', 'N/A')[:30]}...")
            
            # 保存图片
            image_base64 = result.data.get('image_base64')
            if image_base64:
                output_path = OUTPUT_DIR / f"{i:02d}_{case['name']}.png"
                with open(output_path, 'wb') as f:
                    f.write(base64.b64decode(image_base64))
                print(f"    保存: {output_path.name}")
        else:
            print(f"    ✗ 失败: {result.error}")
    
    # 测试带真实图片
    if test_image_path.exists():
        print(f"\n[7] 测试: 带真实图片")
        
        result = await engine.execute({
            "category": "景点",
            "name": "广东小新疆",
            "description": "藏在广东的绝美秘境，比伊犁更有性价比，草原湖泊应有尽有",
            "price": "免费",
            "location": "广东清远",
            "features": "草原风光, 湖泊倒影, 日落绝美, 适合露营",
            "target_audience": "自驾游",
            "image_path": str(test_image_path),
        })
        
        if result.success:
            content = result.data.get('generated_content', {})
            print(f"    ✓ 成功!")
            print(f"    布局: {result.data.get('layout')}, 主题: {result.data.get('theme')}")
            print(f"    标题: {content.get('title', 'N/A')}")
            
            output_path = OUTPUT_DIR / "07_带真实图片_广东小新疆.png"
            with open(output_path, 'wb') as f:
                f.write(base64.b64decode(result.data.get('image_base64')))
            print(f"    保存: {output_path.name}")
        else:
            print(f"    ✗ 失败: {result.error}")
    
    print("\n" + "=" * 60)
    print(f"✓ 完成! 输出目录: {OUTPUT_DIR}")
    print("=" * 60)


async def test_all_layouts():
    """测试所有布局 (强制指定布局)"""
    print("\n" + "=" * 60)
    print("测试所有布局 (强制指定)")
    print("=" * 60)
    
    engine = PosterSmartEngine()
    
    # 基础内容
    base_params = {
        "category": "景点",
        "name": "西湖十景",
        "description": "杭州最美的风景线",
        "price": "免费",
        "location": "杭州",
        "features": "湖光山色, 历史古迹, 文化底蕴",
        "skip_ai": True,  # 跳过 AI，直接用备用生成
    }
    
    for i, layout in enumerate(ALL_LAYOUTS, 1):
        theme = ALL_THEMES[i % len(ALL_THEMES)]
        print(f"\n[{i}] 布局: {layout}, 主题: {theme}")
        
        params = {**base_params, "override_layout": layout, "override_theme": theme}
        result = await engine.execute(params)
        
        if result.success:
            print(f"    ✓ 成功!")
            output_path = OUTPUT_DIR / f"layout_{i:02d}_{layout}.png"
            with open(output_path, 'wb') as f:
                f.write(base64.b64decode(result.data.get('image_base64')))
            print(f"    保存: {output_path.name}")
        else:
            print(f"    ✗ 失败: {result.error}")
    
    print("\n✓ 所有布局测试完成!")


if __name__ == "__main__":
    asyncio.run(test_smart_poster())
    asyncio.run(test_all_layouts())