TravelContentCreator/poster_v2/service.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
海报生成服务 V2
提供统一的海报生成接口，供 AIGC 引擎调用
"""

import io
import base64
import logging
from typing import Dict, Any, Optional, List
from PIL import Image

from .factory import PosterFactory
from .schemas.content import PosterContent
from .schemas.theme import get_theme

logger = logging.getLogger(__name__)


class PosterServiceV2:
    """
    海报生成服务 V2
    
    功能:
    1. 接收内容参数，生成海报
    2. 支持 URL/路径/Base64 图片输入
    3. 输出 Base64 编码的图片
    """
    
    def __init__(self):
        self.factory = PosterFactory()
    
    def generate(
        self,
        layout: str = "hero_bottom",
        theme: str = "ocean",
        content: Dict[str, Any] = None,
        image_base64: str = None,
        image_path: str = None,
        output_format: str = "PNG"
    ) -> Dict[str, Any]:
        """
        生成海报
        
        Args:
            layout: 布局名称 (hero_bottom, overlay_center, etc.)
            theme: 主题名称 (ocean, sunset, etc.)
            content: 内容字典
            image_base64: 背景图片 Base64 (可选)
            image_path: 背景图片路径 (可选)
            output_format: 输出格式 (PNG/JPEG)
            
        Returns:
            {
                "success": True,
                "image_base64": "...",
                "layout": "hero_bottom",
                "theme": "ocean"
            }
        """
        try:
            content = content or {}
            
            # 加载背景图片
            bg_image = None
            if image_base64:
                bg_image = self._load_image_from_base64(image_base64)
            elif image_path:
                bg_image = self._load_image_from_path(image_path)
            
            if bg_image:
                content['image'] = bg_image
            
            # 创建 PosterContent
            poster_content = PosterContent.from_dict(content)
            
            # 生成海报
            poster_image = self.factory.generate_from_content(
                poster_content,
                layout=layout,
                theme=theme
            )
            
            # 转换为 Base64
            image_base64_output = self._image_to_base64(poster_image, output_format)
            
            return {
                "success": True,
                "image_base64": image_base64_output,
                "layout": layout,
                "theme": theme,
            }
            
        except Exception as e:
            logger.error(f"海报生成失败: {e}", exc_info=True)
            return {
                "success": False,
                "error": str(e)
            }
    
    def generate_batch(
        self,
        items: List[Dict[str, Any]],
        default_layout: str = None,
        default_theme: str = None
    ) -> List[Dict[str, Any]]:
        """
        批量生成海报
        
        Args:
            items: 内容列表，每项可包含 layout, theme, content, image_base64
            default_layout: 默认布局
            default_theme: 默认主题
            
        Returns:
            结果列表
        """
        results = []
        for item in items:
            result = self.generate(
                layout=item.get('layout', default_layout or 'hero_bottom'),
                theme=item.get('theme', default_theme or 'ocean'),
                content=item.get('content', {}),
                image_base64=item.get('image_base64'),
                image_path=item.get('image_path'),
            )
            results.append(result)
        return results
    
    def suggest_layout(self, content: Dict[str, Any]) -> str:
        """根据内容推荐布局"""
        return self.factory.suggest_layout(content)
    
    def suggest_theme(self, category: str = None) -> str:
        """根据分类推荐主题"""
        return self.factory.suggest_theme(category)
    
    def list_layouts(self) -> List[str]:
        """列出所有布局"""
        return self.factory.list_layouts()
    
    def list_themes(self) -> List[str]:
        """列出所有主题"""
        return self.factory.list_themes()
    
    def _load_image_from_base64(self, image_base64: str) -> Image.Image:
        """从 Base64 加载图片"""
        # 移除 data URL 前缀
        if ',' in image_base64:
            image_base64 = image_base64.split(',')[1]
        
        image_bytes = base64.b64decode(image_base64)
        return Image.open(io.BytesIO(image_bytes))
    
    def _load_image_from_path(self, image_path: str) -> Image.Image:
        """从路径加载图片"""
        return Image.open(image_path)
    
    def _image_to_base64(self, image: Image.Image, format: str = "PNG") -> str:
        """图片转 Base64"""
        buffer = io.BytesIO()
        
        # 直接转换为 RGB，保留颜色（不做透明合成）
        if image.mode == "RGBA":
            image = image.convert("RGB")
        elif image.mode != "RGB":
            image = image.convert("RGB")
        
        image.save(buffer, format=format.upper())
        return base64.b64encode(buffer.getvalue()).decode('utf-8')


# 单例服务
_service_instance = None


def get_poster_service_v2() -> PosterServiceV2:
    """获取海报服务实例"""
    global _service_instance
    if _service_instance is None:
        _service_instance = PosterServiceV2()
    return _service_instance