88 lines
3.3 KiB
Python
88 lines
3.3 KiB
Python
#!/usr/bin/env python
|
|
# -*- coding: utf-8 -*-
|
|
|
|
import os
|
|
import sys
|
|
import json
|
|
import logging
|
|
import argparse
|
|
from datetime import datetime
|
|
|
|
# 添加项目根目录到路径
|
|
sys.path.append(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 main():
|
|
parser = argparse.ArgumentParser(description='测试额外配图选择功能')
|
|
parser.add_argument('--run_id', type=str, required=True, help='指定运行ID')
|
|
parser.add_argument('--topic_index', type=int, required=True, help='指定主题索引')
|
|
parser.add_argument('--variant_index', type=int, required=True, help='指定变体索引')
|
|
parser.add_argument('--source_dir', type=str, required=True, help='源图像目录路径')
|
|
parser.add_argument('--output_dir', type=str, default='./result', help='输出目录路径')
|
|
parser.add_argument('--num_images', type=int, default=3, help='要选择的配图数量')
|
|
args = parser.parse_args()
|
|
|
|
# 检查源目录是否存在
|
|
if not os.path.exists(args.source_dir) or not os.path.isdir(args.source_dir):
|
|
logger.error(f"源图像目录不存在: {args.source_dir}")
|
|
return 1
|
|
|
|
# 检查输出目录结构
|
|
run_dir = os.path.join(args.output_dir, args.run_id)
|
|
variant_dir = os.path.join(run_dir, f"{args.topic_index}_{args.variant_index}")
|
|
poster_dir = os.path.join(variant_dir, "poster")
|
|
|
|
if not os.path.exists(poster_dir):
|
|
logger.error(f"海报目录不存在: {poster_dir}")
|
|
return 1
|
|
|
|
# 查找海报元数据文件
|
|
metadata_files = [f for f in os.listdir(poster_dir)
|
|
if f.endswith("_metadata.json") and os.path.isfile(os.path.join(poster_dir, f))]
|
|
|
|
if not metadata_files:
|
|
logger.error(f"未找到海报元数据文件")
|
|
return 1
|
|
|
|
metadata_path = os.path.join(poster_dir, metadata_files[0])
|
|
logger.info(f"使用元数据文件: {metadata_path}")
|
|
|
|
# 初始化输出处理器
|
|
output_handler = FileSystemOutputHandler(args.output_dir)
|
|
|
|
# 调用额外配图选择函数
|
|
logger.info(f"开始为运行ID {args.run_id} 主题 {args.topic_index} 变体 {args.variant_index} 选择额外配图...")
|
|
image_paths = select_additional_images(
|
|
run_id=args.run_id,
|
|
topic_index=args.topic_index,
|
|
variant_index=args.variant_index,
|
|
poster_metadata_path=metadata_path,
|
|
source_image_dir=args.source_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, args.run_id, f'{args.topic_index}_{args.variant_index}', 'image')}")
|
|
logger.info("测试成功完成!")
|
|
else:
|
|
logger.error("未能选择任何额外配图")
|
|
return 1
|
|
|
|
return 0
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main()) |