TravelContentCreator/poster/utils.py

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

"""
海报生成工具类
提供图像处理、文本渲染和颜色提取等功能
"""

import os
import logging
import random
from typing import Tuple, List, Dict, Any, Optional, Union
from pathlib import Path
import colorsys
import math
from PIL import Image, ImageDraw, ImageFont, ImageFilter, ImageEnhance

logger = logging.getLogger(__name__)


class ImageProcessor:
    """图像处理工具类"""
    
    def __init__(self):
        self.logger = logging.getLogger(__name__)
    
    def resize_image(self, image: Image.Image, target_size: Tuple[int, int], 
                    keep_aspect: bool = True, crop: bool = False) -> Image.Image:
        """
        调整图像大小
        
        Args:
            image: 原始图像
            target_size: 目标尺寸 (宽, 高)
            keep_aspect: 是否保持宽高比
            crop: 是否裁剪以适应目标尺寸
            
        Returns:
            调整后的图像
        """
        if not keep_aspect:
            return image.resize(target_size, Image.LANCZOS)
        
        # 保持宽高比
        target_width, target_height = target_size
        width, height = image.size
        
        # 计算缩放比例
        ratio_w = target_width / width
        ratio_h = target_height / height
        
        if crop:
            # 裁剪模式：先缩放到较大的尺寸，然后裁剪中心部分
            ratio = max(ratio_w, ratio_h)
            new_width = int(width * ratio)
            new_height = int(height * ratio)
            resized = image.resize((new_width, new_height), Image.LANCZOS)
            
            # 计算裁剪区域
            left = (new_width - target_width) // 2
            top = (new_height - target_height) // 2
            right = left + target_width
            bottom = top + target_height
            
            return resized.crop((left, top, right, bottom))
        else:
            # 非裁剪模式：缩放到适应目标尺寸的最大尺寸
            ratio = min(ratio_w, ratio_h)
            new_width = int(width * ratio)
            new_height = int(height * ratio)
            
            return image.resize((new_width, new_height), Image.LANCZOS)
    
    def apply_overlay(self, image: Image.Image, color: Tuple[int, int, int], 
                     opacity: float = 0.5) -> Image.Image:
        """
        在图像上应用半透明覆盖层
        
        Args:
            image: 原始图像
            color: 覆盖层颜色 (R, G, B)
            opacity: 不透明度 (0.0-1.0)
            
        Returns:
            应用覆盖层后的图像
        """
        overlay = Image.new('RGBA', image.size, color + (int(255 * opacity),))
        
        if image.mode != 'RGBA':
            image = image.convert('RGBA')
            
        return Image.alpha_composite(image, overlay)
    
    def apply_blur(self, image: Image.Image, radius: int = 5) -> Image.Image:
        """
        应用模糊效果
        
        Args:
            image: 原始图像
            radius: 模糊半径
            
        Returns:
            模糊后的图像
        """
        return image.filter(ImageFilter.GaussianBlur(radius))
    
    def create_rounded_corners(self, image: Image.Image, radius: int = 20) -> Image.Image:
        """
        创建圆角图像
        
        Args:
            image: 原始图像
            radius: 圆角半径
            
        Returns:
            圆角处理后的图像
        """
        # 创建一个透明的画布
        rounded = Image.new('RGBA', image.size, (0, 0, 0, 0))
        draw = ImageDraw.Draw(rounded)
        
        # 绘制圆角矩形
        width, height = image.size
        draw.rounded_rectangle([(0, 0), (width, height)], radius, fill=(255, 255, 255, 255))
        
        # 确保图像是RGBA模式
        if image.mode != 'RGBA':
            image = image.convert('RGBA')
            
        # 使用圆角矩形作为蒙版
        result = Image.new('RGBA', image.size, (0, 0, 0, 0))
        result.paste(image, (0, 0), mask=rounded)
        
        return result
    
    def add_shadow(self, image: Image.Image, offset: Tuple[int, int] = (5, 5), 
                  radius: int = 10, color: Tuple[int, int, int] = (0, 0, 0), 
                  opacity: float = 0.7) -> Image.Image:
        """
        为图像添加阴影效果
        
        Args:
            image: 原始图像
            offset: 阴影偏移量 (x, y)
            radius: 阴影模糊半径
            color: 阴影颜色
            opacity: 阴影不透明度
            
        Returns:
            添加阴影后的图像
        """
        # 确保图像是RGBA模式
        if image.mode != 'RGBA':
            image = image.convert('RGBA')
            
        # 创建阴影蒙版
        shadow = Image.new('RGBA', image.size, (0, 0, 0, 0))
        shadow_draw = ImageDraw.Draw(shadow)
        shadow_draw.rectangle([(0, 0), image.size], fill=color + (int(255 * opacity),))
        
        # 应用模糊
        shadow = shadow.filter(ImageFilter.GaussianBlur(radius))
        
        # 创建结果图像
        result = Image.new('RGBA', (
            image.width + abs(offset[0]),
            image.height + abs(offset[1])
        ), (0, 0, 0, 0))
        
        # 放置阴影和原图
        x_offset = max(0, offset[0])
        y_offset = max(0, offset[1])
        result.paste(shadow, (x_offset, y_offset))
        result.paste(image, (max(0, -offset[0]), max(0, -offset[1])), mask=image)
        
        return result


class TextRenderer:
    """文本渲染工具类"""
    
    def __init__(self, font_dir: str = "assets/font"):
        self.font_dir = font_dir
        self.logger = logging.getLogger(__name__)
        self.default_font = self._load_default_font()
        
    def _load_default_font(self, size: int = 24) -> Optional[ImageFont.FreeTypeFont]:
        """加载默认字体"""
        try:
            # 尝试加载系统字体
            system_fonts = [
                "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",  # Linux
                "/System/Library/Fonts/STHeiti Light.ttc",  # macOS
                "C:\\Windows\\Fonts\\msyh.ttc"  # Windows
            ]
            
            # 首先尝试加载指定目录中的字体
            if os.path.exists(self.font_dir):
                font_files = [f for f in os.listdir(self.font_dir) 
                            if f.endswith(('.ttf', '.otf', '.ttc'))]
                if font_files:
                    return ImageFont.truetype(os.path.join(self.font_dir, font_files[0]), size)
            
            # 然后尝试系统字体
            for font_path in system_fonts:
                if os.path.exists(font_path):
                    return ImageFont.truetype(font_path, size)
                    
            # 最后使用默认字体
            return ImageFont.load_default()
            
        except Exception as e:
            self.logger.warning(f"加载默认字体失败: {e}")
            return ImageFont.load_default()

    def draw_text_with_outline(self, draw: ImageDraw.Draw, 
                             position: Tuple[int, int],
                             text: str, 
                             font: ImageFont.FreeTypeFont,
                             text_color: Tuple[int, int, int, int] = (255, 255, 255, 255),
                             outline_color: Tuple[int, int, int, int] = (0, 0, 0, 255),
                             outline_width: int = 2):
        """
        绘制带描边的文字
        
        Args:
            draw: PIL绘图对象
            position: 文字位置
            text: 文字内容
            font: 字体对象
            text_color: 文字颜色
            outline_color: 描边颜色
            outline_width: 描边宽度
        """
        x, y = position
        
        # 绘制描边
        for offset_x in range(-outline_width, outline_width + 1):
            for offset_y in range(-outline_width, outline_width + 1):
                if offset_x == 0 and offset_y == 0:
                    continue
                draw.text((x + offset_x, y + offset_y), text, font=font, fill=outline_color)
        
        # 绘制文字
        draw.text(position, text, font=font, fill=text_color)
    def draw_rounded_rectangle(self, draw: ImageDraw.Draw,
                             position: Tuple[int, int],
                             size: Tuple[int, int],
                             radius: int,
                             fill_color: Tuple[int, int, int, int],
                             outline_color: Optional[Tuple[int, int, int, int]] = None,
                             outline_width: int = 0):
        """
        绘制圆角矩形
        
        Args:
            draw: PIL绘图对象
            position: 左上角位置
            size: 矩形大小
            radius: 圆角半径
            fill_color: 填充颜色
            outline_color: 边框颜色
            outline_width: 边框宽度
        """
        x, y = position
        width, height = size
        
        # 确保尺寸有效
        if width <= 0 or height <= 0:
            return
        
        # 限制圆角半径
        radius = min(radius, width // 2, height // 2)
        
        # 创建圆角矩形路径
        # 这是一个简化版本，PIL的较新版本有更好的圆角矩形支持
        if radius > 0:
            # 绘制中心矩形
            draw.rectangle([x + radius, y, x + width - radius, y + height], fill=fill_color)
            draw.rectangle([x, y + radius, x + width, y + height - radius], fill=fill_color)
            
            # 绘制四个圆角
            draw.pieslice([x, y, x + 2*radius, y + 2*radius], 180, 270, fill=fill_color)
            draw.pieslice([x + width - 2*radius, y, x + width, y + 2*radius], 270, 360, fill=fill_color)
            draw.pieslice([x, y + height - 2*radius, x + 2*radius, y + height], 90, 180, fill=fill_color)
            draw.pieslice([x + width - 2*radius, y + height - 2*radius, x + width, y + height], 0, 90, fill=fill_color)
        else:
            # 普通矩形
            draw.rectangle([x, y, x + width, y + height], fill=fill_color)
        
        # 绘制边框（如果需要）
        if outline_color and outline_width > 0:
            # 简化的边框绘制 - 使用线条而不是矩形避免坐标错误
            for i in range(outline_width):
                offset = i
                # 确保坐标有效
                if radius > 0:
                    # 上边
                    if x + radius + offset < x + width - radius - offset:
                        draw.line([x + radius + offset, y + offset, 
                                  x + width - radius - offset, y + offset], 
                                 fill=outline_color, width=1)
                    # 下边
                    if x + radius + offset < x + width - radius - offset and y + height - offset >= y + offset:
                        draw.line([x + radius + offset, y + height - offset, 
                                  x + width - radius - offset, y + height - offset], 
                                 fill=outline_color, width=1)
                    # 左边
                    if y + radius + offset < y + height - radius - offset:
                        draw.line([x + offset, y + radius + offset, 
                                  x + offset, y + height - radius - offset], 
                                 fill=outline_color, width=1)
                    # 右边
                    if y + radius + offset < y + height - radius - offset:
                        draw.line([x + width - offset, y + radius + offset, 
                                  x + width - offset, y + height - radius - offset], 
                                 fill=outline_color, width=1)
                else:
                    # 普通矩形边框
                    draw.rectangle([x + offset, y + offset, x + width - offset, y + height - offset], 
                                  outline=outline_color, width=1)
                                  
    def draw_text_with_shadow(self, draw: ImageDraw.Draw,
                            position: Tuple[int, int],
                            text: str,
                            font: ImageFont.FreeTypeFont,
                            text_color: Tuple[int, int, int, int] = (255, 255, 255, 255),
                            shadow_color: Tuple[int, int, int, int] = (0, 0, 0, 128),
                            shadow_offset: Tuple[int, int] = (2, 2)):
        """
        绘制带阴影的文字
        
        Args:
            draw: PIL绘图对象
            position: 文字位置
            text: 文字内容
            font: 字体对象
            text_color: 文字颜色
            shadow_color: 阴影颜色
            shadow_offset: 阴影偏移
        """
        x, y = position
        shadow_x, shadow_y = shadow_offset
        
        # 绘制阴影
        draw.text((x + shadow_x, y + shadow_y), text, font=font, fill=shadow_color)
        
        # 绘制文字
        draw.text(position, text, font=font, fill=text_color)
    
    def get_font(self, font_name: Optional[str] = None, size: int = 24) -> ImageFont.FreeTypeFont:
        """
        获取指定字体
        
        Args:
            font_name: 字体文件名，如果为None则使用默认字体
            size: 字体大小
            
        Returns:
            字体对象
        """
        if not font_name:
            return self.default_font.font_variant(size=size)
        
        try:
            # 尝试从字体目录加载
            font_path = os.path.join(self.font_dir, font_name)
            if os.path.exists(font_path):
                return ImageFont.truetype(font_path, size)
            
            # 尝试作为绝对路径加载
            if os.path.exists(font_name):
                return ImageFont.truetype(font_name, size)
                
            self.logger.warning(f"找不到字体: {font_name}，使用默认字体")
            return self.default_font.font_variant(size=size)
            
        except Exception as e:
            self.logger.warning(f"加载字体 {font_name} 失败: {e}")
            return self.default_font.font_variant(size=size)
    
    def draw_text(self, image: Image.Image, text: str, position: Tuple[int, int], 
                 font_name: Optional[str] = None, font_size: int = 24, 
                 color: Tuple[int, int, int] = (0, 0, 0), 
                 align: str = "left", max_width: Optional[int] = None) -> Image.Image:
        """
        在图像上绘制文本
        
        Args:
            image: 目标图像
            text: 要绘制的文本
            position: 文本位置 (x, y)
            font_name: 字体文件名
            font_size: 字体大小
            color: 文本颜色 (R, G, B)
            align: 对齐方式 ('left', 'center', 'right')
            max_width: 最大宽度，超过会自动换行
            
        Returns:
            绘制文本后的图像
        """
        draw = ImageDraw.Draw(image)
        font = self.get_font(font_name, font_size)
        
        if not max_width:
            # 简单绘制文本
            x, y = position
            
            if align == "center":
                text_width, _ = draw.textsize(text, font=font)
                x -= text_width // 2
            elif align == "right":
                text_width, _ = draw.textsize(text, font=font)
                x -= text_width
                
            draw.text((x, y), text, font=font, fill=color)
            return image
        
        # 处理自动换行
        lines = []
        current_line = ""
        
        for word in text.split():
            test_line = current_line + word + " "
            text_width, _ = draw.textsize(test_line, font=font)
            
            if text_width <= max_width:
                current_line = test_line
            else:
                lines.append(current_line)
                current_line = word + " "
                
        if current_line:
            lines.append(current_line)
            
        # 绘制每一行
        x, y = position
        line_height = font_size * 1.2  # 估计行高
        
        for i, line in enumerate(lines):
            line_y = y + i * line_height
            
            if align == "center":
                line_width, _ = draw.textsize(line, font=font)
                line_x = x - line_width // 2
            elif align == "right":
                line_width, _ = draw.textsize(line, font=font)
                line_x = x - line_width
            else:
                line_x = x
                
            draw.text((line_x, line_y), line, font=font, fill=color)
            
        return image

    def calculate_optimal_font_size(self, text: str, 
                                target_width: int, 
                                font_name: Optional[str] = None,
                                max_size: int = 120, 
                                min_size: int = 10) -> int:
        """
        计算最适合的字体大小
        
        Args:
            text: 文字内容
            target_width: 目标宽度
            font_name: 字体文件名
            max_size: 最大字体大小
            min_size: 最小字体大小
            
        Returns:
            最适合的字体大小
        """
        if not text.strip():
            return min_size
    
    
        # 二分查找最佳字体大小
        left, right = min_size, max_size
        best_size = min_size
        
        while left <= right:
            mid_size = (left + right) // 2
            
            try:
                if font_path:
                    font = self._load_default_font(mid_size)
                else:
                    font = ImageFont.load_default()
                
                # 获取文字边界框
                bbox = font.getbbox(text)
                text_width = bbox[2] - bbox[0]
                
                if text_width <= target_width:
                    best_size = mid_size
                    left = mid_size + 1
                else:
                    right = mid_size - 1
                    
            except Exception:
                right = mid_size - 1
        
        return best_size
    
    def get_text_size(self, text: str, font: ImageFont.FreeTypeFont) -> Tuple[int, int]:
        """
        获取文字的尺寸
        
        Args:
            text: 文字内容
            font: 字体对象
            
        Returns:
            文字尺寸 (width, height)
        """
        bbox = font.getbbox(text)
        return bbox[2] - bbox[0], bbox[3] - bbox[1]
    
    def draw_multiline_text(self, image: Image.Image, text: str, position: Tuple[int, int],
                          font_name: Optional[str] = None, font_size: int = 24,
                          color: Tuple[int, int, int] = (0, 0, 0),
                          align: str = "left", spacing: int = 4,
                          max_width: Optional[int] = None) -> Image.Image:
        """
        绘制多行文本，支持自动换行
        
        Args:
            image: 目标图像
            text: 要绘制的多行文本
            position: 文本起始位置 (x, y)
            font_name: 字体文件名
            font_size: 字体大小
            color: 文本颜色 (R, G, B)
            align: 对齐方式 ('left', 'center', 'right')
            spacing: 行间距
            max_width: 最大宽度，超过会自动换行
            
        Returns:
            绘制文本后的图像
        """
        draw = ImageDraw.Draw(image)
        font = self.get_font(font_name, font_size)
        
        # 处理预先分行的文本
        paragraphs = text.split('\n')
        all_lines = []
        
        for paragraph in paragraphs:
            if not max_width:
                all_lines.append(paragraph)
                continue
                
            # 自动换行处理
            words = paragraph.split()
            current_line = ""
            
            for word in words:
                test_line = current_line + word + " "
                text_width, _ = draw.textsize(test_line, font=font)
                
                if text_width <= max_width:
                    current_line = test_line
                else:
                    all_lines.append(current_line)
                    current_line = word + " "
                    
            if current_line:
                all_lines.append(current_line)
        
        # 绘制所有行
        x, y = position
        line_height = font_size + spacing
        
        for i, line in enumerate(all_lines):
            line_y = y + i * line_height
            
            if align == "center":
                line_width, _ = draw.textsize(line, font=font)
                line_x = x - line_width // 2
            elif align == "right":
                line_width, _ = draw.textsize(line, font=font)
                line_x = x - line_width
            else:
                line_x = x
                
            draw.text((line_x, line_y), line, font=font, fill=color)
            
        return image


class ColorExtractor:
    """颜色提取工具类"""
    
    def __init__(self):
        self.logger = logging.getLogger(__name__)
    
    def extract_dominant_color(self, image: Image.Image, num_colors: int = 5) -> List[Tuple[int, int, int]]:
        """
        提取图像中的主要颜色
        
        Args:
            image: 输入图像
            num_colors: 要提取的颜色数量
            
        Returns:
            颜色列表，按主导程度排序
        """
        # 缩小图像以加快处理速度
        small_image = image.resize((100, 100))
        
        # 确保图像是RGB模式
        if small_image.mode != 'RGB':
            small_image = small_image.convert('RGB')
            
        # 获取所有像素
        pixels = list(small_image.getdata())
        
        # 计算颜色频率
        color_counts = {}
        for pixel in pixels:
            if pixel in color_counts:
                color_counts[pixel] += 1
            else:
                color_counts[pixel] = 1
                
        # 按频率排序
        sorted_colors = sorted(color_counts.items(), key=lambda x: x[1], reverse=True)
        
        # 返回前N个颜色
        return [color for color, _ in sorted_colors[:num_colors]]
    
    def extract_color_palette(self, image: Image.Image, num_colors: int = 5) -> List[Tuple[int, int, int]]:
        """
        提取图像的颜色调色板
        使用k-means聚类算法的简化版本
        
        Args:
            image: 输入图像
            num_colors: 调色板中的颜色数量
            
        Returns:
            颜色列表
        """
        # 缩小图像以加快处理速度
        small_image = image.resize((100, 100))
        
        # 确保图像是RGB模式
        if small_image.mode != 'RGB':
            small_image = small_image.convert('RGB')
            
        # 获取所有像素
        pixels = list(small_image.getdata())
        
        # 随机选择初始中心点
        centers = random.sample(pixels, num_colors)
        
        # 简单的k-means聚类 (最多迭代10次)
        for _ in range(10):
            clusters = [[] for _ in range(num_colors)]
            
            # 将每个像素分配给最近的中心点
            for pixel in pixels:
                distances = [self._color_distance(pixel, center) for center in centers]
                closest_center = distances.index(min(distances))
                clusters[closest_center].append(pixel)
                
            # 更新中心点
            new_centers = []
            for cluster in clusters:
                if not cluster:
                    # 如果聚类为空，保持原中心点
                    new_centers.append(centers[len(new_centers)])
                    continue
                    
                # 计算聚类中所有像素的平均值
                r_sum = sum(pixel[0] for pixel in cluster)
                g_sum = sum(pixel[1] for pixel in cluster)
                b_sum = sum(pixel[2] for pixel in cluster)
                
                new_center = (
                    r_sum // len(cluster),
                    g_sum // len(cluster),
                    b_sum // len(cluster)
                )
                new_centers.append(new_center)
                
            # 检查是否收敛
            if centers == new_centers:
                break
                
            centers = new_centers
            
        return centers
    
    def _color_distance(self, color1: Tuple[int, int, int], color2: Tuple[int, int, int]) -> float:
        """计算两个颜色之间的欧几里得距离"""
        r1, g1, b1 = color1
        r2, g2, b2 = color2
        return math.sqrt((r1 - r2) ** 2 + (g1 - g2) ** 2 + (b1 - b2) ** 2)
    
    def generate_complementary_color(self, color: Tuple[int, int, int]) -> Tuple[int, int, int]:
        """
        生成互补色
        
        Args:
            color: 输入颜色 (R, G, B)
            
        Returns:
            互补色 (R, G, B)
        """
        r, g, b = color
        
        # 转换为HSV
        h, s, v = colorsys.rgb_to_hsv(r / 255, g / 255, b / 255)
        
        # 互补色的色相相差180度
        h = (h + 0.5) % 1.0
        
        # 转回RGB
        r, g, b = colorsys.hsv_to_rgb(h, s, v)
        
        return (int(r * 255), int(g * 255), int(b * 255))
    
    def generate_color_scheme(self, base_color: Tuple[int, int, int], 
                            scheme_type: str = "complementary", 
                            num_colors: int = 5) -> List[Tuple[int, int, int]]:
        """
        生成配色方案
        
        Args:
            base_color: 基础颜色 (R, G, B)
            scheme_type: 配色方案类型 ('complementary', 'analogous', 'triadic', 'monochromatic')
            num_colors: 生成的颜色数量
            
        Returns:
            颜色列表
        """
        r, g, b = base_color
        h, s, v = colorsys.rgb_to_hsv(r / 255, g / 255, b / 255)
        colors = []
        
        if scheme_type == "complementary":
            # 互补色方案
            colors.append(base_color)
            h_comp = (h + 0.5) % 1.0
            r_comp, g_comp, b_comp = colorsys.hsv_to_rgb(h_comp, s, v)
            colors.append((int(r_comp * 255), int(g_comp * 255), int(b_comp * 255)))
            
            # 添加额外的颜色
            for i in range(num_colors - 2):
                h_extra = (h + (i + 1) * 0.1) % 1.0
                r_extra, g_extra, b_extra = colorsys.hsv_to_rgb(h_extra, s, v)
                colors.append((int(r_extra * 255), int(g_extra * 255), int(b_extra * 255)))
                
        elif scheme_type == "analogous":
            # 类似色方案
            for i in range(num_colors):
                h_analog = (h + (i - num_colors // 2) * 0.05) % 1.0
                r_analog, g_analog, b_analog = colorsys.hsv_to_rgb(h_analog, s, v)
                colors.append((int(r_analog * 255), int(g_analog * 255), int(b_analog * 255)))
                
        elif scheme_type == "triadic":
            # 三色方案
            for i in range(3):
                h_triad = (h + i / 3) % 1.0
                r_triad, g_triad, b_triad = colorsys.hsv_to_rgb(h_triad, s, v)
                colors.append((int(r_triad * 255), int(g_triad * 255), int(b_triad * 255)))
                
            # 添加额外的颜色
            for i in range(num_colors - 3):
                h_extra = (h + (i + 1) * 0.1) % 1.0
                r_extra, g_extra, b_extra = colorsys.hsv_to_rgb(h_extra, s, v)
                colors.append((int(r_extra * 255), int(g_extra * 255), int(b_extra * 255)))
                
        elif scheme_type == "monochromatic":
            # 单色方案
            for i in range(num_colors):
                s_mono = max(0.1, min(1.0, s - 0.3 + i * 0.6 / (num_colors - 1)))
                v_mono = max(0.2, min(1.0, v - 0.3 + i * 0.6 / (num_colors - 1)))
                r_mono, g_mono, b_mono = colorsys.hsv_to_rgb(h, s_mono, v_mono)
                colors.append((int(r_mono * 255), int(g_mono * 255), int(b_mono * 255)))
                
        else:
            # 默认返回随机颜色
            colors.append(base_color)
            for _ in range(num_colors - 1):
                h_rand = random.random()
                s_rand = random.uniform(0.5, 1.0)
                v_rand = random.uniform(0.5, 1.0)
                r_rand, g_rand, b_rand = colorsys.hsv_to_rgb(h_rand, s_rand, v_rand)
                colors.append((int(r_rand * 255), int(g_rand * 255), int(b_rand * 255)))
                
        return colors