121 lines
4.0 KiB
Python
121 lines
4.0 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
|
||
"""
|
||
海报生成器模块
|
||
"""
|
||
|
||
import os
|
||
import logging
|
||
import random
|
||
from typing import List, Dict, Any, Optional, Union, Tuple
|
||
from pathlib import Path
|
||
from PIL import Image
|
||
|
||
from core.config import ConfigManager, PosterConfig
|
||
from utils.file_io import OutputManager
|
||
from .templates.base_template import BaseTemplate
|
||
from .templates.vibrant_template import VibrantTemplate
|
||
from .templates.business_template import BusinessTemplate
|
||
from .templates.collage_template import CollageTemplate
|
||
from .utils import ImageProcessor
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class PosterGenerator:
|
||
"""
|
||
海报生成器
|
||
负责根据文本内容和配置生成海报图像
|
||
"""
|
||
|
||
def __init__(self, config_manager: ConfigManager, output_manager: OutputManager):
|
||
"""
|
||
初始化海报生成器
|
||
|
||
Args:
|
||
config_manager: 配置管理器
|
||
output_manager: 输出管理器
|
||
"""
|
||
self.config_manager = config_manager
|
||
self.config = config_manager.get_config('poster', PosterConfig)
|
||
self.output_manager = output_manager
|
||
self.image_processor = ImageProcessor()
|
||
|
||
# 注册可用的模板
|
||
self.templates = {
|
||
"vibrant": VibrantTemplate,
|
||
"business": BusinessTemplate,
|
||
"collage": CollageTemplate,
|
||
}
|
||
|
||
def generate_poster(self, content: Dict[str, Any], topic_index: Union[int, str],
|
||
template_name: Optional[str] = None) -> Optional[str]:
|
||
"""
|
||
为指定内容生成海报
|
||
|
||
Args:
|
||
content: 内容数据，包含标题、正文等
|
||
topic_index: 主题索引，用于文件命名
|
||
template_name: 模板名称，如果为None则根据配置选择
|
||
|
||
Returns:
|
||
生成的海报文件路径，如果生成失败则返回None
|
||
"""
|
||
logger.info(f"开始为主题 {topic_index} 生成海报...")
|
||
|
||
# 1. 选择模板
|
||
if not template_name:
|
||
template_name = self._select_template()
|
||
|
||
if template_name not in self.templates:
|
||
logger.error(f"模板 '{template_name}' 不存在")
|
||
return None
|
||
|
||
template_class = self.templates[template_name]
|
||
template = template_class(self.config.target_size)
|
||
|
||
# 2. 准备内容
|
||
title = content.get('title', '')
|
||
text_content = content.get('content', '')
|
||
|
||
if not title or not text_content:
|
||
logger.error("内容缺少标题或正文")
|
||
return None
|
||
|
||
# 3. 生成海报
|
||
try:
|
||
poster_image = template.generate(title, text_content)
|
||
|
||
# 4. 保存海报
|
||
output_dir = self.output_manager.get_topic_dir(topic_index)
|
||
file_name = f"poster_{template_name}.png"
|
||
file_path = output_dir / file_name
|
||
|
||
poster_image.save(file_path)
|
||
logger.info(f"海报已保存到: {file_path}")
|
||
|
||
return str(file_path)
|
||
except Exception as e:
|
||
logger.error(f"生成海报时发生错误: {e}", exc_info=True)
|
||
return None
|
||
|
||
def _select_template(self) -> str:
|
||
"""根据配置选择模板"""
|
||
selection_method = self.config.template_selection
|
||
available_templates = self.config.available_templates
|
||
|
||
# 过滤出实际可用的模板
|
||
valid_templates = [t for t in available_templates if t in self.templates]
|
||
|
||
if not valid_templates:
|
||
logger.warning("没有可用的模板，使用默认模板")
|
||
return next(iter(self.templates.keys()))
|
||
|
||
if selection_method == "random":
|
||
return random.choice(valid_templates)
|
||
elif selection_method in valid_templates:
|
||
return selection_method
|
||
else:
|
||
logger.warning(f"无效的模板选择方法 '{selection_method}'，使用随机选择")
|
||
return random.choice(valid_templates) |