129 lines
4.2 KiB
Python
129 lines
4.2 KiB
Python
#!/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 FileSystemOutputHandler
|
|
from utils.poster_notes_creator import select_additional_images
|
|
|
|
# 配置日志
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
def create_sample_metadata(image_dir, num_images=2):
|
|
"""创建示例海报元数据
|
|
|
|
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_images', 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 = FileSystemOutputHandler(args.output_dir)
|
|
|
|
# 创建测试运行ID和主题/变体索引
|
|
run_id = datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
topic_index = 1
|
|
variant_index = 1
|
|
|
|
# 创建一个模拟的海报元数据文件
|
|
metadata = create_sample_metadata(args.image_dir, 2)
|
|
metadata_path = os.path.join(args.output_dir, f"{run_id}_{topic_index}_{variant_index}_metadata.json")
|
|
|
|
# 保存元数据到文件
|
|
with open(metadata_path, 'w', encoding='utf-8') as f:
|
|
json.dump(metadata, f, ensure_ascii=False, indent=2)
|
|
logger.info(f"模拟海报元数据已保存: {metadata_path}")
|
|
|
|
# 打印将要使用的"已用"图像
|
|
logger.info(f"模拟海报使用的图像: {', '.join(metadata['collage_images'])}")
|
|
|
|
# 调用额外配图选择函数
|
|
logger.info(f"开始选择额外配图...")
|
|
image_paths = select_additional_images(
|
|
run_id=run_id,
|
|
topic_index=topic_index,
|
|
variant_index=variant_index,
|
|
poster_metadata_path=metadata_path,
|
|
source_image_dir=args.image_dir,
|
|
num_additional_images=args.num_images,
|
|
output_handler=output_handler
|
|
)
|
|
|
|
# 打印结果
|
|
if image_paths:
|
|
logger.info(f"已选择并保存 {len(image_paths)} 张额外配图:")
|
|
for i, path in enumerate(image_paths):
|
|
logger.info(f" {i+1}. {path}")
|
|
|
|
logger.info(f"额外配图已保存到: {os.path.join(args.output_dir, run_id, f'{topic_index}_{variant_index}', 'image')}")
|
|
logger.info(f"测试成功完成!")
|
|
else:
|
|
logger.error("未能选择任何额外配图.")
|
|
return 1
|
|
|
|
return 0
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main()) |