232 lines
8.4 KiB
Python
232 lines
8.4 KiB
Python
#!/usr/bin/env python
|
|
# -*- coding: utf-8 -*-
|
|
|
|
"""
|
|
测试PosterGenerator组件
|
|
|
|
该脚本演示如何独立使用PosterGenerator组件创建海报。
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import argparse
|
|
import json
|
|
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 poster_gen
|
|
|
|
def parse_arguments():
|
|
"""解析命令行参数"""
|
|
parser = argparse.ArgumentParser(description='测试海报生成功能')
|
|
parser.add_argument('--base_dir', '-b', type=str,
|
|
help='海报素材基础目录路径，包含贴纸/字体/边框等')
|
|
parser.add_argument('--image', '-i', type=str,
|
|
help='底图路径')
|
|
parser.add_argument('--output_dir', '-o', type=str, default='poster_output',
|
|
help='输出海报保存目录，默认为"poster_output"')
|
|
parser.add_argument('--config', '-c', type=str,
|
|
help='海报配置文件路径')
|
|
return parser.parse_args()
|
|
|
|
def load_config(config_path="poster_gen_config.json"):
|
|
"""加载配置文件"""
|
|
if not os.path.exists(config_path):
|
|
config_path = os.path.join(project_root, config_path)
|
|
if not os.path.exists(config_path):
|
|
raise FileNotFoundError(f"配置文件未找到: {config_path}")
|
|
|
|
with open(config_path, 'r', encoding='utf-8') as f:
|
|
config = json.load(f)
|
|
return config
|
|
|
|
def create_test_image(output_path, size=(900, 1200)):
|
|
"""创建测试底图"""
|
|
# 创建一个渐变背景
|
|
img = Image.new('RGB', size, (0, 0, 0))
|
|
pixels = img.load()
|
|
|
|
width, height = size
|
|
for x in range(width):
|
|
for y in range(height):
|
|
r = int(255 * x / width)
|
|
g = int(255 * y / height)
|
|
b = int(127 * (x + y) / (width + height))
|
|
pixels[x, y] = (r, g, b)
|
|
|
|
# 保存图片
|
|
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
|
img.save(output_path)
|
|
print(f"创建测试底图: {output_path}")
|
|
return output_path
|
|
|
|
def sample_text_data():
|
|
"""生成示例文本数据"""
|
|
# 几组示例文本数据，便于测试不同文案效果
|
|
samples = [
|
|
{
|
|
"title": "收藏！泰宁古城",
|
|
"subtitle": "",
|
|
"additional_texts": [
|
|
{"text": "周末游｜必去", "position": "bottom", "size_factor": 0.5},
|
|
{"text": "千年古韵", "position": "bottom", "size_factor": 0.5}
|
|
]
|
|
},
|
|
{
|
|
"title": "拿捏龙门石窟",
|
|
"subtitle": "",
|
|
"additional_texts": [
|
|
{"text": "亲子研学", "position": "bottom", "size_factor": 0.5},
|
|
{"text": "文化瑰宝", "position": "bottom", "size_factor": 0.5}
|
|
]
|
|
},
|
|
{
|
|
"title": "厦门鼓浪屿",
|
|
"subtitle": "音乐之岛",
|
|
"additional_texts": [
|
|
{"text": "出片｜必打卡", "position": "bottom", "size_factor": 0.5}
|
|
]
|
|
}
|
|
]
|
|
return random.choice(samples)
|
|
|
|
def main():
|
|
"""主函数"""
|
|
args = parse_arguments()
|
|
|
|
# 确保输出目录存在
|
|
if not os.path.exists(args.output_dir):
|
|
os.makedirs(args.output_dir)
|
|
|
|
try:
|
|
# 尝试加载系统配置
|
|
try:
|
|
config = load_config()
|
|
print("系统配置加载成功")
|
|
except Exception as e:
|
|
print(f"系统配置加载失败，将使用默认值和命令行参数: {e}")
|
|
config = {}
|
|
|
|
# 确定海报素材基础目录
|
|
base_dir = args.base_dir
|
|
if not base_dir:
|
|
base_dir = config.get("poster_assets_base_dir")
|
|
if not base_dir:
|
|
print("警告: 未提供海报素材基础目录，将使用默认路径")
|
|
base_dir = "poster_assets"
|
|
|
|
print(f"使用海报素材基础目录: {base_dir}")
|
|
|
|
# 初始化PosterGenerator
|
|
print("初始化PosterGenerator...")
|
|
poster_generator = poster_gen.PosterGenerator(base_dir=base_dir)
|
|
|
|
# 确定底图路径
|
|
image_path = args.image
|
|
if not image_path or not os.path.exists(image_path):
|
|
# 尝试从配置中获取示例图片
|
|
try:
|
|
img_base_dir = config.get("image_base_dir")
|
|
object_name = "测试景点"
|
|
img_subdir = config.get("modify_image_subdir", "modify")
|
|
|
|
if img_base_dir and os.path.exists(img_base_dir):
|
|
example_dir = os.path.join(img_base_dir, img_subdir)
|
|
if os.path.exists(example_dir):
|
|
for root, dirs, files in os.walk(example_dir):
|
|
for file in files:
|
|
if file.lower().endswith(('.jpg', '.jpeg', '.png')):
|
|
image_path = os.path.join(root, file)
|
|
break
|
|
if image_path:
|
|
break
|
|
except Exception:
|
|
pass
|
|
|
|
# 如果仍未找到图片，创建测试图片
|
|
if not image_path or not os.path.exists(image_path):
|
|
print("未找到有效底图，创建测试图片...")
|
|
image_path = create_test_image(os.path.join(args.output_dir, "test_base_image.jpg"))
|
|
|
|
print(f"使用底图: {image_path}")
|
|
|
|
# 加载底图
|
|
try:
|
|
base_image = Image.open(image_path).convert('RGBA')
|
|
print(f"底图加载成功，尺寸: {base_image.size}")
|
|
except Exception as e:
|
|
print(f"底图加载失败，将创建纯色底图: {e}")
|
|
base_image = Image.new('RGBA', (900, 1200), (200, 200, 200, 255))
|
|
|
|
# 确定文本数据
|
|
text_data = None
|
|
if args.config and os.path.exists(args.config):
|
|
try:
|
|
with open(args.config, 'r', encoding='utf-8') as f:
|
|
config_data = json.load(f)
|
|
|
|
# 检查是否是海报配置列表
|
|
if isinstance(config_data, list) and len(config_data) > 0:
|
|
# 使用第一个配置
|
|
poster_config = config_data[0]
|
|
text_data = {
|
|
"title": poster_config.get("main_title", "默认标题"),
|
|
"subtitle": "",
|
|
"additional_texts": []
|
|
}
|
|
|
|
# 添加文本内容
|
|
texts = poster_config.get("texts", [])
|
|
if texts and len(texts) > 0:
|
|
for i, text in enumerate(texts):
|
|
text_data["additional_texts"].append({
|
|
"text": text,
|
|
"position": "bottom",
|
|
"size_factor": 0.5
|
|
})
|
|
|
|
print(f"从配置文件加载文本数据成功: {text_data}")
|
|
except Exception as e:
|
|
print(f"配置文件解析失败: {e}")
|
|
text_data = None
|
|
|
|
# 如果未提供配置或解析失败，使用示例文本
|
|
if not text_data:
|
|
text_data = sample_text_data()
|
|
print(f"使用示例文本数据: {text_data}")
|
|
|
|
# 生成海报
|
|
print("\n开始生成海报...")
|
|
start_time = time.time()
|
|
|
|
# 设置是否使用文本背景
|
|
os.environ['USE_TEXT_BG'] = 'False'
|
|
|
|
# 生成海报
|
|
poster_image = poster_generator.create_poster(base_image, text_data)
|
|
|
|
if poster_image:
|
|
# 保存海报
|
|
output_path = os.path.join(args.output_dir, "generated_poster.jpg")
|
|
poster_image.save(output_path, "JPEG", quality=95)
|
|
|
|
elapsed = time.time() - start_time
|
|
print(f"海报生成成功，耗时: {elapsed:.2f}秒")
|
|
print(f"海报已保存至: {output_path}")
|
|
else:
|
|
print("海报生成失败")
|
|
|
|
except Exception as e:
|
|
print(f"程序执行出错: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
if __name__ == "__main__":
|
|
main() |