TravelContentCreator/examples/test_poster_notes.py

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

import os
import sys
import logging
import argparse
import json
import random
from pathlib import Path
from datetime import datetime
from PIL import Image

# 添加项目根目录到路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from utils.output_handler import OutputHandler
from utils.poster_notes_creator import process_poster_for_notes

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class SimpleOutputHandler(OutputHandler):
    """简单的输出处理器，用于测试目的"""
    
    def __init__(self, output_dir):
        """初始化输出处理器
        
        Args:
            output_dir: 输出目录路径
        """
        self.output_dir = output_dir
        os.makedirs(output_dir, exist_ok=True)
        logger.info(f"已初始化输出处理器，输出目录: {output_dir}")
    
    def handle_generated_image(self, run_id, topic_index, variant_index, 
                              image_type, image_data, output_filename, metadata=None):
        """处理生成的图像
        
        Args:
            run_id: 运行ID
            topic_index: 主题索引
            variant_index: 变体索引
            image_type: 图像类型
            image_data: 图像数据
            output_filename: 输出文件名
            metadata: 图像元数据
            
        Returns:
            str: 保存的图像路径
        """
        # 创建目录结构
        output_dir = ""
        if image_type == 'note':
            # 笔记图像保存在image目录
            output_dir = os.path.join(self.output_dir, f"{run_id}/{topic_index}_{variant_index}/image")
        else:
            # 其他类型图像保存在各自的类型目录
            output_dir = os.path.join(self.output_dir, image_type)
            
        os.makedirs(output_dir, exist_ok=True)
        
        # 保存图像
        image_path = os.path.join(output_dir, output_filename)
        image_data.save(image_path)
        
        # 如果有元数据，保存它
        if metadata:
            metadata_path = os.path.splitext(image_path)[0] + '.json'
            with open(metadata_path, 'w', encoding='utf-8') as f:
                json.dump(metadata, f, ensure_ascii=False, indent=2)
        
        logger.info(f"已保存图像: {image_path}")
        return image_path
    
    def handle_poster_configs(self, run_id, topic_index, variant_index, 
                             poster_configs, output_filename):
        """处理海报配置
        
        Args:
            run_id: 运行ID
            topic_index: 主题索引
            variant_index: 变体索引
            poster_configs: 海报配置
            output_filename: 输出文件名
            
        Returns:
            str: 保存的配置路径
        """
        # 创建目录
        config_dir = os.path.join(self.output_dir, 'configs')
        os.makedirs(config_dir, exist_ok=True)
        
        # 保存配置
        config_path = os.path.join(config_dir, output_filename)
        with open(config_path, 'w', encoding='utf-8') as f:
            json.dump(poster_configs, f, ensure_ascii=False, indent=2)
        
        logger.info(f"已保存配置: {config_path}")
        return config_path

def create_sample_metadata(image_dir, num_images=3):
    """创建示例海报元数据
    
    Args:
        image_dir: 图像目录
        num_images: 使用的图像数量
        
    Returns:
        dict: 海报元数据
    """
    # 获取目录中的图像
    image_extensions = ('.jpg', '.jpeg', '.png', '.bmp')
    available_images = [
        f for f in os.listdir(image_dir)
        if os.path.isfile(os.path.join(image_dir, f)) and
        f.lower().endswith(image_extensions)
    ]
    
    if len(available_images) < num_images:
        logger.warning(f"可用图像数量({len(available_images)})少于请求数量({num_images})，将使用所有可用图像")
        selected_images = available_images
    else:
        selected_images = random.sample(available_images, num_images)
    
    # 创建元数据
    metadata = {
        "title": "示例海报",
        "description": "这是一个用于测试的示例海报",
        "collage_images": selected_images,
        "generation_time": datetime.now().isoformat(),
        "style": "modern"
    }
    
    return metadata

def main():
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='测试海报笔记创建器')
    parser.add_argument('--image_dir', type=str, required=True, 
                        help='包含图像的目录路径')
    parser.add_argument('--output_dir', type=str, default='./output',
                       help='输出目录路径')
    parser.add_argument('--num_notes', type=int, default=3,
                       help='要创建的笔记数量')
    
    args = parser.parse_args()
    
    # 验证图像目录
    if not os.path.exists(args.image_dir) or not os.path.isdir(args.image_dir):
        logger.error(f"图像目录不存在: {args.image_dir}")
        return 1
    
    # 创建输出目录
    os.makedirs(args.output_dir, exist_ok=True)
    
    # 初始化输出处理器
    output_handler = SimpleOutputHandler(args.output_dir)
    
    # 创建一个示例海报和元数据
    run_id = datetime.now().strftime("%Y%m%d_%H%M%S")
    topic_index = 0
    variant_index = 0
    
    # 从图像目录中选择一张图像作为"海报"
    image_extensions = ('.jpg', '.jpeg', '.png', '.bmp')
    available_images = [
        f for f in os.listdir(args.image_dir)
        if os.path.isfile(os.path.join(args.image_dir, f)) and
        f.lower().endswith(image_extensions)
    ]
    
    if not available_images:
        logger.error(f"图像目录中没有找到图像: {args.image_dir}")
        return 1
    
    # 选择第一张图像作为海报
    poster_image_name = available_images[0]
    poster_image_path = os.path.join(args.image_dir, poster_image_name)
    
    # 创建示例元数据
    poster_metadata = create_sample_metadata(args.image_dir, 2)
    
    # 保存海报图像
    poster_image = Image.open(poster_image_path)
    saved_poster_path = output_handler.handle_generated_image(
        run_id, 
        topic_index, 
        variant_index, 
        'poster', 
        poster_image, 
        'sample_poster.jpg', 
        poster_metadata
    )
    
    # 保存海报元数据
    poster_metadata_path = os.path.splitext(saved_poster_path)[0] + '.json'
    
    logger.info(f"已创建示例海报: {saved_poster_path}")
    logger.info(f"元数据路径: {poster_metadata_path}")
    
    # 处理海报笔记
    logger.info("开始创建海报笔记...")
    note_paths = process_poster_for_notes(
        run_id,
        topic_index,
        variant_index,
        saved_poster_path,
        poster_metadata_path,
        args.image_dir,
        args.num_notes,
        output_handler
    )
    
    logger.info(f"创建了 {len(note_paths)} 个笔记图像:")
    for i, path in enumerate(note_paths):
        logger.info(f"  {i+1}. {path}")
    
    return 0

if __name__ == "__main__":
    sys.exit(main())