TravelContentCreator/domain/aigc/shared/image_processor.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
图片处理器
提供图片处理的统一接口
"""

import logging
import base64
from io import BytesIO
from typing import Optional, Tuple, Union, List
from pathlib import Path

logger = logging.getLogger(__name__)

try:
    from PIL import Image
    PIL_AVAILABLE = True
except ImportError:
    PIL_AVAILABLE = False
    logger.warning("PIL 未安装，图片处理功能受限")


class ImageProcessor:
    """
    图片处理器
    
    提供：
    - Base64 编解码
    - 尺寸调整
    - 格式转换
    - 基础图片操作
    """
    
    def __init__(self):
        if not PIL_AVAILABLE:
            logger.warning("PIL 未安装，部分功能不可用")
    
    def decode_base64(self, base64_str: str) -> Optional['Image.Image']:
        """
        解码 Base64 图片
        
        Args:
            base64_str: Base64 编码的图片字符串
            
        Returns:
            PIL Image 对象
        """
        if not PIL_AVAILABLE:
            raise RuntimeError("PIL 未安装")
        
        try:
            # 移除可能的 data URL 前缀
            if ',' in base64_str:
                base64_str = base64_str.split(',', 1)[1]
            
            # 解码
            image_data = base64.b64decode(base64_str)
            image = Image.open(BytesIO(image_data))
            
            # 转换为 RGBA
            if image.mode != 'RGBA':
                image = image.convert('RGBA')
            
            return image
            
        except Exception as e:
            logger.error(f"解码 Base64 图片失败: {e}")
            return None
    
    def to_base64(
        self,
        image: 'Image.Image',
        format: str = 'PNG',
        quality: int = 95
    ) -> str:
        """
        将图片编码为 Base64
        
        Args:
            image: PIL Image 对象
            format: 图片格式 (PNG, JPEG, WEBP)
            quality: 质量 (仅对 JPEG/WEBP 有效)
            
        Returns:
            Base64 编码字符串
        """
        if not PIL_AVAILABLE:
            raise RuntimeError("PIL 未安装")
        
        try:
            buffer = BytesIO()
            
            # 保存到缓冲区
            save_kwargs = {'format': format}
            if format.upper() in ('JPEG', 'WEBP'):
                save_kwargs['quality'] = quality
            
            # JPEG 不支持 RGBA
            if format.upper() == 'JPEG' and image.mode == 'RGBA':
                image = image.convert('RGB')
            
            image.save(buffer, **save_kwargs)
            
            # 编码
            return base64.b64encode(buffer.getvalue()).decode('utf-8')
            
        except Exception as e:
            logger.error(f"编码图片为 Base64 失败: {e}")
            raise
    
    def resize(
        self,
        image: 'Image.Image',
        size: Tuple[int, int],
        keep_aspect: bool = True,
        fill_color: Tuple[int, int, int, int] = (255, 255, 255, 255)
    ) -> 'Image.Image':
        """
        调整图片尺寸
        
        Args:
            image: PIL Image 对象
            size: 目标尺寸 (width, height)
            keep_aspect: 是否保持宽高比
            fill_color: 填充颜色 (RGBA)
            
        Returns:
            调整后的图片
        """
        if not PIL_AVAILABLE:
            raise RuntimeError("PIL 未安装")
        
        if not keep_aspect:
            return image.resize(size, Image.Resampling.LANCZOS)
        
        # 保持宽高比
        original_ratio = image.width / image.height
        target_ratio = size[0] / size[1]
        
        if original_ratio > target_ratio:
            # 原图更宽，按宽度缩放
            new_width = size[0]
            new_height = int(size[0] / original_ratio)
        else:
            # 原图更高，按高度缩放
            new_height = size[1]
            new_width = int(size[1] * original_ratio)
        
        resized = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
        
        # 创建目标尺寸的画布
        result = Image.new('RGBA', size, fill_color)
        
        # 居中粘贴
        x = (size[0] - new_width) // 2
        y = (size[1] - new_height) // 2
        result.paste(resized, (x, y))
        
        return result
    
    def crop_center(
        self,
        image: 'Image.Image',
        size: Tuple[int, int]
    ) -> 'Image.Image':
        """
        中心裁剪
        
        Args:
            image: PIL Image 对象
            size: 目标尺寸 (width, height)
            
        Returns:
            裁剪后的图片
        """
        if not PIL_AVAILABLE:
            raise RuntimeError("PIL 未安装")
        
        # 先缩放到合适大小
        scale = max(size[0] / image.width, size[1] / image.height)
        new_size = (int(image.width * scale), int(image.height * scale))
        image = image.resize(new_size, Image.Resampling.LANCZOS)
        
        # 中心裁剪
        left = (image.width - size[0]) // 2
        top = (image.height - size[1]) // 2
        right = left + size[0]
        bottom = top + size[1]
        
        return image.crop((left, top, right, bottom))
    
    def apply_blur(
        self,
        image: 'Image.Image',
        radius: int = 10
    ) -> 'Image.Image':
        """
        应用模糊效果
        
        Args:
            image: PIL Image 对象
            radius: 模糊半径
            
        Returns:
            模糊后的图片
        """
        if not PIL_AVAILABLE:
            raise RuntimeError("PIL 未安装")
        
        from PIL import ImageFilter
        return image.filter(ImageFilter.GaussianBlur(radius=radius))
    
    def load_from_file(self, file_path: Union[str, Path]) -> Optional['Image.Image']:
        """
        从文件加载图片
        
        Args:
            file_path: 文件路径
            
        Returns:
            PIL Image 对象
        """
        if not PIL_AVAILABLE:
            raise RuntimeError("PIL 未安装")
        
        try:
            image = Image.open(file_path)
            if image.mode != 'RGBA':
                image = image.convert('RGBA')
            return image
        except Exception as e:
            logger.error(f"加载图片失败: {file_path}, {e}")
            return None
    
    def save_to_file(
        self,
        image: 'Image.Image',
        file_path: Union[str, Path],
        format: Optional[str] = None,
        quality: int = 95
    ) -> bool:
        """
        保存图片到文件
        
        Args:
            image: PIL Image 对象
            file_path: 文件路径
            format: 图片格式（不指定则从扩展名推断）
            quality: 质量
            
        Returns:
            是否成功
        """
        if not PIL_AVAILABLE:
            raise RuntimeError("PIL 未安装")
        
        try:
            file_path = Path(file_path)
            file_path.parent.mkdir(parents=True, exist_ok=True)
            
            save_kwargs = {}
            if format:
                save_kwargs['format'] = format
            if file_path.suffix.lower() in ('.jpg', '.jpeg', '.webp'):
                save_kwargs['quality'] = quality
            
            image.save(file_path, **save_kwargs)
            return True
            
        except Exception as e:
            logger.error(f"保存图片失败: {file_path}, {e}")
            return False