TravelContentCreator/examples/test_image_processing.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
测试脚本：专门测试图像处理功能
不依赖于AI接口，仅测试图像拼贴和海报生成功能
"""
import os
import sys
import json
import argparse
import logging
import random
import numpy as np
from datetime import datetime

# 添加项目根目录到Python路径
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if PROJECT_ROOT not in sys.path:
    sys.path.append(PROJECT_ROOT)

# 导入所需的图像处理模块
try:
    import core.simple_collage as simple_collage
    import core.poster_gen as poster_gen
    from utils.resource_loader import ResourceLoader
except ImportError as e:
    logging.critical(f"导入模块失败: {e}")
    logging.critical(f"确保'{PROJECT_ROOT}'在sys.path中，并且所有依赖已安装")
    sys.exit(1)

def load_config(config_path="poster_gen_config.json"):
    """加载配置文件"""
    try:
        with open(config_path, 'r', encoding='utf-8') as f:
            config = json.load(f)
        logging.info(f"从 {config_path} 加载配置成功")
        return config
    except FileNotFoundError:
        logging.error(f"配置文件未找到: {config_path}")
        return None
    except json.JSONDecodeError:
        logging.error(f"配置文件解析失败: {config_path}，请检查JSON格式")
        return None
    except Exception as e:
        logging.exception(f"加载配置文件时发生错误:")
        return None

def create_test_images(output_dir, count=5, size=(800, 600)):
    """创建测试用的彩色图像"""
    os.makedirs(output_dir, exist_ok=True)
    logging.info(f"创建 {count} 张测试图像到 {output_dir}")
    
    image_paths = []
    for i in range(count):
        # 创建随机彩色图像
        img = np.zeros((size[1], size[0], 3), dtype=np.uint8)
        color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
        img[:] = color
        
        # 保存图像
        img_path = os.path.join(output_dir, f"test_img_{i+1}.jpg")
        try:
            import cv2
            cv2.imwrite(img_path, img)
            image_paths.append(img_path)
            logging.info(f"创建测试图像: {img_path}")
        except Exception as e:
            logging.error(f"创建测试图像失败: {e}")
    
    return image_paths

def test_collage(config, output_dir, test_images_dir=None, collage_count=3):
    """测试图像拼贴功能"""
    logging.info("=== 开始测试图像拼贴功能 ===")
    
    # 准备输出目录
    collage_output_dir = os.path.join(output_dir, "collages")
    os.makedirs(collage_output_dir, exist_ok=True)
    
    # 准备测试图像
    if test_images_dir is None or not os.path.exists(test_images_dir):
        test_images_dir = os.path.join(output_dir, "test_images")
        image_paths = create_test_images(test_images_dir)
    else:
        image_paths = [os.path.join(test_images_dir, f) for f in os.listdir(test_images_dir) 
                      if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
    
    if not image_paths:
        logging.error("无有效测试图像可用")
        return False
    
    # 测试不同风格的拼贴
    styles = ["standard", "asymmetric", "filmstrip", "overlap", "mosaic"]
    success_count = 0
    
    for i in range(collage_count):
        style = styles[i % len(styles)]
        try:
            logging.info(f"生成风格为 '{style}' 的拼贴图 {i+1}/{collage_count}")
            output_path = os.path.join(collage_output_dir, f"collage_{style}_{i+1}.jpg")
            
            # 创建拼贴图
            collage_creator = simple_collage.ImageCollageCreator(
                target_size=(900, 900),
                style=style
            )
            
            # 随机选择3-5张图片
            num_images = random.randint(3, min(5, len(image_paths)))
            selected_images = random.sample(image_paths, num_images)
            
            # 加载图像
            images = []
            for img_path in selected_images:
                try:
                    import cv2
                    img = cv2.imread(img_path)
                    if img is not None:
                        images.append(img)
                except Exception as e:
                    logging.warning(f"加载图像 {img_path} 失败: {e}")
            
            if not images:
                logging.warning(f"无法加载任何图像，跳过拼贴 {i+1}")
                continue
            
            # 创建拼贴
            result = collage_creator.create_collage(images)
            
            # 保存结果
            cv2.imwrite(output_path, result)
            logging.info(f"拼贴图保存到: {output_path}")
            success_count += 1
            
        except Exception as e:
            logging.exception(f"创建拼贴图 {i+1} 时出错:")
    
    logging.info(f"拼贴测试完成: 成功 {success_count}/{collage_count}")
    return success_count > 0

def test_poster_generation(config, output_dir, collage_dir=None):
    """测试海报生成功能"""
    logging.info("=== 开始测试海报生成功能 ===")
    
    # 准备输出目录
    poster_output_dir = os.path.join(output_dir, "posters")
    os.makedirs(poster_output_dir, exist_ok=True)
    
    # 获取海报素材目录
    poster_assets_dir = config.get("poster_assets_base_dir")
    if not poster_assets_dir or not os.path.exists(poster_assets_dir):
        logging.error(f"海报素材目录不存在或未配置: {poster_assets_dir}")
        return False
    
    # 获取或创建底图
    base_images = []
    
    # 使用之前创建的拼贴图
    if collage_dir and os.path.exists(collage_dir):
        collage_files = [f for f in os.listdir(collage_dir) if f.lower().endswith('.jpg')]
        for file in collage_files:
            base_images.append(os.path.join(collage_dir, file))
    
    # 如果没有拼贴图或数量不足，创建纯色测试图像
    if len(base_images) < 3:
        test_img_dir = os.path.join(output_dir, "base_images")
        new_images = create_test_images(test_img_dir, count=3, size=(900, 1200))
        base_images.extend(new_images)
    
    # 准备一些测试文本
    test_texts = [
        {
            "title": "云南大理古城悠闲一日游",
            "subtitle": "探寻苍山洱海间的千年历史",
            "content": "大理古城，一座被苍山洱海环抱的千年古城，漫步其中，感受白族建筑与现代艺术的完美融合。"
        },
        {
            "title": "秋日北京香山红叶之旅",
            "subtitle": "首都最美的秋天风景线",
            "content": "每年深秋，香山红叶如火如荼，是摄影爱好者的天堂，也是周末休闲的绝佳去处。"
        },
        {
            "title": "青岛海滨度假攻略",
            "subtitle": "啤酒与海鲜的夏日盛宴",
            "content": '青岛，被誉为"东方瑞士"，在这里，你可以感受到欧式建筑与中国沿海城市的独特魅力。'
        }
    ]
    
    # 创建海报生成器
    try:
        poster_generator = poster_gen.PosterGenerator(
            poster_save_dir=poster_output_dir,
            assets_base_dir=poster_assets_dir,
            poster_size=tuple(config.get("poster_target_size", [900, 1200]))
        )
        logging.info("海报生成器初始化成功")
    except Exception as e:
        logging.exception("初始化海报生成器失败:")
        return False
    
    # 生成多张测试海报
    success_count = 0
    for i, base_img_path in enumerate(base_images[:3]):
        try:
            text_data = test_texts[i % len(test_texts)]
            output_path = os.path.join(poster_output_dir, f"poster_{i+1}.jpg")
            
            logging.info(f"生成海报 {i+1}/3，使用底图: {os.path.basename(base_img_path)}")
            
            # 调用海报生成
            poster_generator.create_poster(
                base_image_path=base_img_path,
                poster_text_config={
                    "title": text_data["title"],
                    "subtitle": text_data["subtitle"],
                    "content": text_data["content"],
                    "has_second_paragraph": random.random() < 0.5
                },
                output_path=output_path
            )
            
            if os.path.exists(output_path):
                logging.info(f"海报已保存到: {output_path}")
                success_count += 1
            else:
                logging.warning(f"海报创建似乎成功，但文件未找到: {output_path}")
                
        except Exception as e:
            logging.exception(f"创建海报 {i+1} 时出错:")
    
    logging.info(f"海报测试完成: 成功 {success_count}/3")
    return success_count > 0

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='poster_gen_config.json', help='配置文件路径')
    parser.add_argument('--output', type=str, default=None, help='输出目录，默认为result/test_image_TIMESTAMP')
    parser.add_argument('--collage_only', action='store_true', help='仅测试拼贴功能')
    parser.add_argument('--poster_only', action='store_true', help='仅测试海报生成功能')
    args = parser.parse_args()
    
    # 加载配置
    logging.info(f"加载配置: {args.config}")
    config = load_config(args.config)
    if not config:
        logging.critical("配置加载失败，无法继续测试")
        sys.exit(1)
    
    # 准备输出目录
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    if args.output:
        output_base_dir = args.output
    else:
        output_base_dir = os.path.join(config.get("output_dir", "result"), f"test_image_{timestamp}")
    
    os.makedirs(output_base_dir, exist_ok=True)
    logging.info(f"测试输出将保存到: {output_base_dir}")
    
    # 执行测试
    test_results = []
    
    # 如果两个only标志都没设置，则都测试
    test_all = not args.collage_only and not args.poster_only
    
    # 测试拼贴
    if args.collage_only or test_all:
        collage_success = test_collage(config, output_base_dir)
        test_results.append(("拼贴功能", collage_success))
        
        # 如果只测试拼贴或需要连续测试，且拼贴成功了
        if collage_success and (args.collage_only or test_all):
            collage_dir = os.path.join(output_base_dir, "collages")
        else:
            collage_dir = None
    else:
        collage_dir = None
    
    # 测试海报生成
    if args.poster_only or test_all:
        poster_success = test_poster_generation(config, output_base_dir, collage_dir)
        test_results.append(("海报生成功能", poster_success))
    
    # 输出结果摘要
    logging.info("\n=== 测试结果摘要 ===")
    all_success = True
    for name, result in test_results:
        status = "成功" if result else "失败"
        logging.info(f"{name}: {status}")
        all_success = all_success and result
    
    logging.info(f"\n测试{'全部通过' if all_success else '部分失败'}!")
    logging.info(f"所有测试输出保存在: {output_base_dir}")
    
    sys.exit(0 if all_success else 1)

if __name__ == "__main__":
    main()