TravelContentCreator/examples/test_simple_collage.py

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

"""
测试图片拼贴功能

该脚本演示如何独立使用simple_collage模块创建图片拼贴。
"""

import os
import sys
import argparse
import time
import random
from pathlib import Path
from PIL import Image

# 将项目根目录添加到PATH
project_root = str(Path(__file__).parent.parent.absolute())
if project_root not in sys.path:
    sys.path.append(project_root)

from core import simple_collage

def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='测试图片拼贴功能')
    parser.add_argument('--input_dir', '-i', type=str, 
                        help='输入图片目录路径')
    parser.add_argument('--output_dir', '-o', type=str, default='collage_output',
                        help='输出拼贴图保存目录，默认为"collage_output"')
    parser.add_argument('--width', '-w', type=int, default=900,
                        help='拼贴图宽度，默认为900')
    parser.add_argument('--height', '-h', type=int, default=1200,
                        help='拼贴图高度，默认为1200')
    parser.add_argument('--count', '-c', type=int, default=3,
                        help='生成拼贴图数量，默认为3')
    parser.add_argument('--seed', '-s', type=int, default=None,
                        help='随机种子，用于重现结果')
    return parser.parse_args()

def get_image_files(directory):
    """获取目录中的所有图片文件"""
    if not os.path.exists(directory):
        raise FileNotFoundError(f"目录不存在: {directory}")
    
    image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.webp', '.tiff']
    image_files = []
    
    for root, _, files in os.walk(directory):
        for file in files:
            if any(file.lower().endswith(ext) for ext in image_extensions):
                image_files.append(os.path.join(root, file))
    
    if not image_files:
        raise ValueError(f"目录中没有找到有效的图片文件: {directory}")
    
    return image_files

def create_test_images(directory, count=5, size=(640, 480)):
    """如果没有提供输入目录，创建测试图片"""
    if not os.path.exists(directory):
        os.makedirs(directory)
    
    colors = [
        (255, 0, 0),     # 红色
        (0, 255, 0),     # 绿色
        (0, 0, 255),     # 蓝色
        (255, 255, 0),   # 黄色
        (255, 0, 255),   # 紫色
        (0, 255, 255),   # 青色
        (255, 165, 0),   # 橙色
        (128, 0, 128),   # 紫色
        (210, 105, 30),  # 巧克力色
        (0, 128, 128)    # 墨绿色
    ]
    
    image_files = []
    for i in range(count):
        # 随机选择颜色
        color = colors[i % len(colors)]
        
        # 创建纯色图片
        img = Image.new('RGB', size, color)
        
        # 保存图片
        file_path = os.path.join(directory, f"test_image_{i+1}.jpg")
        img.save(file_path)
        image_files.append(file_path)
        print(f"创建测试图片: {file_path}")
    
    return image_files

def main():
    """主函数"""
    args = parse_arguments()
    
    # 设置随机种子
    if args.seed is not None:
        random.seed(args.seed)
    
    # 确保输出目录存在
    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)
    
    # 获取输入图片
    if not args.input_dir:
        print("未提供输入目录，创建测试图片...")
        test_img_dir = os.path.join(args.output_dir, "test_images")
        image_files = create_test_images(test_img_dir, count=7)
        input_dir = test_img_dir
    else:
        input_dir = args.input_dir
        try:
            # 尝试获取图片文件列表，验证目录有效
            get_image_files(input_dir)
        except Exception as e:
            print(f"错误: {e}")
            return
    
    # 设置目标尺寸
    target_size = (args.width, args.height)
    print(f"使用输入目录: {input_dir}")
    print(f"拼贴图尺寸: {target_size}")
    
    # 生成拼贴图
    print(f"开始生成{args.count}张拼贴图...")
    start_time = time.time()
    
    try:
        # 创建与测试不同样式的拼贴图
        collage_creator = simple_collage.ImageCollageCreator()
        
        # 1. 处理整个目录，获取多个拼贴图
        print("方法1: 使用process_directory批量生成")
        collages, used_image_names = simple_collage.process_directory(
            input_dir,
            target_size=target_size,
            output_count=args.count
        )
        
        if not collages:
            print("拼贴图生成失败！")
            return
        
        # 保存拼贴图
        for i, collage in enumerate(collages):
            output_path = os.path.join(args.output_dir, f"collage_auto_{i+1}.png")
            collage.save(output_path)
            # 输出使用的图片名称
            if i < len(used_image_names):
                print(f"拼贴图已保存: {output_path}，使用图片: {used_image_names[i]}")
            else:
                print(f"拼贴图已保存: {output_path}")
        
        # 2. 使用不同风格创建拼贴图
        print("\n方法2: 测试不同风格")
        styles = ["grid", "asymmetrical", "filmstrip", "overlap", "mosaic"]
        
        for style in styles:
            print(f"创建 {style} 风格拼贴图...")
            try:
                collage, selected_images = collage_creator.create_collage_with_style(input_dir, style, target_size)
                if collage:
                    output_path = os.path.join(args.output_dir, f"collage_style_{style}.png")
                    collage.save(output_path)
                    print(f"风格拼贴图已保存: {output_path}，使用图片: {selected_images}")
                else:
                    print(f"创建 {style} 风格失败: 未返回有效拼贴图")
            except Exception as e:
                print(f"创建 {style} 风格失败: {e}")
        
        elapsed = time.time() - start_time
        print(f"\n拼贴图生成完成，总耗时: {elapsed:.2f}秒")
        print(f"输出目录: {os.path.abspath(args.output_dir)}")
    
    except Exception as e:
        print(f"拼贴图生成出错: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()