追加了海报测试程序
This commit is contained in:
parent
c81a3080ca
commit
ed8d005203
202
scripts/test_poster_gen.py
Normal file
202
scripts/test_poster_gen.py
Normal file
@ -0,0 +1,202 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding: utf-8 -*-
|
||||
|
||||
import os
|
||||
import sys
|
||||
import argparse
|
||||
from PIL import Image
|
||||
import random
|
||||
import json
|
||||
from datetime import datetime
|
||||
|
||||
# 导入海报生成模块
|
||||
from core.poster_gen import PosterGenerator, PosterConfig
|
||||
|
||||
def parse_arguments():
|
||||
"""解析命令行参数"""
|
||||
parser = argparse.ArgumentParser(description='测试海报生成模块')
|
||||
|
||||
# 必需参数
|
||||
parser.add_argument('--assets-dir', type=str, required=True,
|
||||
help='海报资源目录，包含frames, font, stickers等子目录')
|
||||
|
||||
# 可选参数
|
||||
parser.add_argument('--image', type=str, default=None,
|
||||
help='输入图片路径（单张图片测试）')
|
||||
parser.add_argument('--image-dir', type=str, default=None,
|
||||
help='输入图片目录（批量测试）')
|
||||
parser.add_argument('--config', type=str, default=None,
|
||||
help='海报配置文件路径')
|
||||
parser.add_argument('--output-dir', type=str, default='poster_test_output',
|
||||
help='输出目录')
|
||||
parser.add_argument('--count', type=int, default=1,
|
||||
help='生成海报数量')
|
||||
parser.add_argument('--frame-possibility', type=float, default=0.7,
|
||||
help='添加边框的概率')
|
||||
parser.add_argument('--text-bg-possibility', type=float, default=0.0,
|
||||
help='添加文本背景的概率')
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
def create_sample_config():
|
||||
"""创建示例配置"""
|
||||
return [
|
||||
{
|
||||
"index": 0,
|
||||
"main_title": "美丽的海滩度假胜地",
|
||||
"texts": [
|
||||
"阳光沙滩，椰林树影，尽享热带风情",
|
||||
"给自己一次难忘的海岛旅行"
|
||||
]
|
||||
},
|
||||
{
|
||||
"index": 1,
|
||||
"main_title": "壮丽的山川河流",
|
||||
"texts": [
|
||||
"峰峦叠嶂，溪水潺潺，领略自然之美",
|
||||
"放慢脚步，感受大自然的馈赠"
|
||||
]
|
||||
},
|
||||
{
|
||||
"index": 2,
|
||||
"main_title": "古老的城市文化",
|
||||
"texts": [
|
||||
"千年古迹，历史遗存，品味人文气息",
|
||||
"穿越时空，探寻文明的足迹"
|
||||
]
|
||||
}
|
||||
]
|
||||
|
||||
def load_config(config_path):
|
||||
"""加载海报配置"""
|
||||
if not config_path:
|
||||
return create_sample_config()
|
||||
|
||||
try:
|
||||
if os.path.exists(config_path):
|
||||
with open(config_path, 'r', encoding='utf-8') as f:
|
||||
return json.load(f)
|
||||
else:
|
||||
print(f"配置文件不存在: {config_path}，使用默认配置")
|
||||
return create_sample_config()
|
||||
except Exception as e:
|
||||
print(f"加载配置失败: {e}，使用默认配置")
|
||||
return create_sample_config()
|
||||
|
||||
def get_image_files(image_dir):
|
||||
"""获取目录中的图片文件"""
|
||||
if not os.path.exists(image_dir):
|
||||
print(f"图片目录不存在: {image_dir}")
|
||||
return []
|
||||
|
||||
extensions = ('.jpg', '.jpeg', '.png')
|
||||
return [os.path.join(image_dir, f) for f in os.listdir(image_dir)
|
||||
if f.lower().endswith(extensions)]
|
||||
|
||||
def main():
|
||||
# 解析命令行参数
|
||||
args = parse_arguments()
|
||||
|
||||
# 确保资源目录存在
|
||||
if not os.path.exists(args.assets_dir):
|
||||
print(f"错误: 海报资源目录不存在: {args.assets_dir}")
|
||||
return
|
||||
|
||||
# 创建输出目录
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
|
||||
# 加载配置
|
||||
configs = load_config(args.config)
|
||||
print(f"加载了 {len(configs)} 个海报配置")
|
||||
|
||||
# 创建海报生成器
|
||||
poster_generator = PosterGenerator(args.assets_dir, args.output_dir)
|
||||
poster_generator.set_img_frame_possibility(args.frame_possibility)
|
||||
poster_generator.set_text_bg_possibility(args.text_bg_possibility)
|
||||
|
||||
# 环境设置
|
||||
os.environ['USE_TEXT_BG'] = 'True' if args.text_bg_possibility > 0 else 'False'
|
||||
|
||||
# 获取图片列表
|
||||
image_list = []
|
||||
if args.image:
|
||||
if os.path.exists(args.image):
|
||||
image_list = [args.image]
|
||||
else:
|
||||
print(f"错误: 指定的图片不存在: {args.image}")
|
||||
return
|
||||
elif args.image_dir:
|
||||
image_list = get_image_files(args.image_dir)
|
||||
if not image_list:
|
||||
print(f"错误: 在目录中没有找到图片: {args.image_dir}")
|
||||
return
|
||||
else:
|
||||
print("错误: 必须提供单张图片路径或图片目录")
|
||||
return
|
||||
|
||||
print(f"找到 {len(image_list)} 张图片用于测试")
|
||||
|
||||
# 生成海报
|
||||
generated_count = 0
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
|
||||
for i in range(args.count):
|
||||
# 选择图片
|
||||
if not image_list:
|
||||
print("错误: 没有可用的图片")
|
||||
break
|
||||
|
||||
image_path = random.choice(image_list)
|
||||
|
||||
# 选择配置
|
||||
config = random.choice(configs)
|
||||
|
||||
# 构建文本数据
|
||||
text_data = {
|
||||
"title": config["main_title"],
|
||||
"subtitle": "",
|
||||
"additional_texts": []
|
||||
}
|
||||
|
||||
# 添加额外文本
|
||||
if "texts" in config and config["texts"]:
|
||||
for idx, text in enumerate(config["texts"]):
|
||||
if text:
|
||||
position = "bottom" if idx > 0 else "middle"
|
||||
text_data["additional_texts"].append({
|
||||
"text": text,
|
||||
"position": position,
|
||||
"size_factor": 0.6 if idx > 0 else 0.8
|
||||
})
|
||||
|
||||
try:
|
||||
# 加载图片
|
||||
print(f"\n生成海报 {i+1}/{args.count}")
|
||||
print(f"使用图片: {image_path}")
|
||||
print(f"标题: {text_data['title']}")
|
||||
|
||||
# 尝试先加载PIL图像对象
|
||||
img = Image.open(image_path).convert('RGBA')
|
||||
|
||||
# 生成海报
|
||||
poster = poster_generator.create_poster(img, text_data)
|
||||
|
||||
if poster:
|
||||
# 保存结果
|
||||
output_filename = f"poster_{timestamp}_{i+1}.jpg"
|
||||
output_path = os.path.join(args.output_dir, output_filename)
|
||||
poster.save(output_path)
|
||||
print(f"成功生成海报: {output_path}")
|
||||
generated_count += 1
|
||||
else:
|
||||
print(f"生成海报失败")
|
||||
except Exception as e:
|
||||
print(f"处理图片时出错: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
|
||||
print(f"\n测试完成! 成功生成 {generated_count}/{args.count} 张海报")
|
||||
print(f"输出目录: {args.output_dir}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Loading…
x
Reference in New Issue
Block a user