TravelContentCreator/examples/test_image_variations.py

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

import os
import sys
import logging
import argparse
from pathlib import Path
from PIL import Image

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

from utils.poster_notes_creator import PosterNotesCreator
from utils.output_handler import FileSystemOutputHandler

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

def main():
    parser = argparse.ArgumentParser(description='测试图像变化效果')
    parser.add_argument('--input_image', type=str, required=True, help='输入图像路径')
    parser.add_argument('--output_dir', type=str, default='./output/image_variations', help='输出目录')
    parser.add_argument('--variations', type=int, default=5, help='要生成的变体数量')
    parser.add_argument('--strength', type=str, default='medium', choices=['low', 'medium', 'high'], 
                        help='变化强度: low(低), medium(中), high(高)')
    parser.add_argument('--aspect_ratio', type=str, default='3:4', help='目标宽高比，格式如"3:4"')
    parser.add_argument('--extra_effects', action='store_true', default=True, 
                        help='是否应用额外效果（噪点、锐化等）')
    args = parser.parse_args()

    # 解析宽高比
    try:
        aspect_parts = args.aspect_ratio.split(':')
        if len(aspect_parts) != 2:
            raise ValueError("宽高比格式无效")
        target_ratio = (int(aspect_parts[0]), int(aspect_parts[1]))
    except Exception as e:
        logger.error(f"解析宽高比时出错: {e}")
        logger.info("使用默认宽高比 3:4")
        target_ratio = (3, 4)

    # 检查输入图像
    if not os.path.exists(args.input_image) or not os.path.isfile(args.input_image):
        logger.error(f"输入图像不存在: {args.input_image}")
        return 1
    
    # 创建输出目录
    os.makedirs(args.output_dir, exist_ok=True)
    
    # 加载输入图像
    try:
        original_image = Image.open(args.input_image)
        logger.info(f"已加载图像: {args.input_image}, 尺寸: {original_image.size}")
    except Exception as e:
        logger.error(f"加载图像时出错: {e}")
        return 1
    
    # 初始化处理器
    output_handler = FileSystemOutputHandler(os.path.dirname(args.output_dir))
    creator = PosterNotesCreator(output_handler)
    
    # 保存原始图像信息
    width, height = original_image.size
    aspect = width / height
    logger.info(f"原始图像宽高比: {aspect:.4f} ({width}x{height})")
    
    # 先保存原始图像，调整为目标比例但不添加变化
    try:
        original_processed = creator.process_image_to_aspect_ratio(
            original_image, 
            target_ratio, 
            add_variation=False
        )
        original_output_path = os.path.join(args.output_dir, "original_processed.jpg")
        original_processed.save(original_output_path)
        processed_width, processed_height = original_processed.size
        logger.info(f"已保存原始处理图像（无变化）: {original_output_path}")
        logger.info(f"处理后尺寸: {processed_width}x{processed_height}, 宽高比: {processed_width/processed_height:.4f}")
    except Exception as e:
        logger.error(f"处理原始图像时出错: {e}")
    
    # 生成多个变体
    logger.info(f"开始生成 {args.variations} 个图像变体（强度: {args.strength}, 额外效果: {'开启' if args.extra_effects else '关闭'}）...")
    
    for i in range(args.variations):
        try:
            # 处理图像并添加变化
            varied_image = creator.process_image_to_aspect_ratio(
                original_image, 
                target_ratio, 
                add_variation=True,
                seed=i+1,  # 使用不同的种子
                variation_strength=args.strength,
                extra_effects=args.extra_effects
            )
            
            # 保存变体
            output_path = os.path.join(args.output_dir, f"variation_{i+1}_{args.strength}.jpg")
            varied_image.save(output_path)
            logger.info(f"已保存变体 {i+1}: {output_path}")
            
        except Exception as e:
            logger.error(f"生成变体 {i+1} 时出错: {e}")
    
    # 如果有额外效果，也生成一组无额外效果的变体来对比
    if args.extra_effects and args.strength != 'low':
        logger.info("生成无额外效果的变体用于对比...")
        for i in range(min(3, args.variations)):  # 只生成少量用于对比
            try:
                varied_image = creator.process_image_to_aspect_ratio(
                    original_image, 
                    target_ratio, 
                    add_variation=True,
                    seed=i+100,  # 使用不同的种子
                    variation_strength=args.strength,
                    extra_effects=False  # 关闭额外效果
                )
                
                output_path = os.path.join(args.output_dir, f"no_extra_effects_{i+1}.jpg")
                varied_image.save(output_path)
                logger.info(f"已保存无额外效果变体 {i+1}: {output_path}")
                
            except Exception as e:
                logger.error(f"生成无额外效果变体 {i+1} 时出错: {e}")
    
    logger.info(f"图像变体生成完成。所有图像都保存在: {args.output_dir}")
    
    # 列出各种配置对照
    logger.info("\n图像变化配置参考:")
    logger.info("- 低强度变化: 亮度/对比度/饱和度 ±3%, 最大旋转 ±0.5°")
    logger.info("- 中强度变化: 亮度/对比度/饱和度 ±5%, 最大旋转 ±1.0°")
    logger.info("- 高强度变化: 亮度/对比度/饱和度 ±8%, 最大旋转 ±2.0°")
    logger.info("- 额外效果包括: 微弱噪点, 色相调整, 边缘处理, 锐化/模糊")
    
    return 0

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