PictureEdit/image_manager.py

"""
作者：太一
微信：taiyi1224
邮箱：shuobo1224@qq.com
"""

import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageOps


class ImageManager:
    """图像管理工具类，处理图像加载、变换和保存"""

    @staticmethod
    def load_image(file_path):
        """加载图像并转换为RGBA格式"""
        try:
            return Image.open(file_path).convert("RGBA")
        except Exception as e:
            print(f"加载图像失败: {e}")
            return None

    @staticmethod
    def apply_mask(image, shape):
        """应用遮罩到图像"""
        if shape == "rect":
            return image  # 矩形不需要遮罩

        # 创建与图像相同大小的透明遮罩
        mask = Image.new('L', image.size, 0)
        draw = ImageDraw.Draw(mask)

        # 根据形状绘制遮罩
        if shape == "round":
            # 圆形遮罩
            draw.ellipse((0, 0, image.size[0], image.size[1]), fill=255)
        elif shape == "heart":
            # 心形遮罩（简化版）
            x, y = image.size
            points = [
                (x // 2, y // 5),
                (x // 5, y // 2),
                (x // 5, y * 3 // 4),
                (x // 2, y * 4 // 5),
                (x * 4 // 5, y * 3 // 4),
                (x * 4 // 5, y // 2),
                (x // 2, y // 5)
            ]
            draw.polygon(points, fill=255)
        elif shape == "star":
            # 星形遮罩（简化版）
            x, y = image.size
            center_x, center_y = x // 2, y // 2
            radius = min(center_x, center_y)
            points = []
            for i in range(10):
                angle = 0.1 * i * 3.14159
                r = radius if i % 2 == 0 else radius * 0.4
                points.append((
                    center_x + r * (1 if i % 4 < 2 else -1) * abs(angle % (2 * 3.14159) - 3.14159 / 2) ** 0.5,
                    center_y + r * (1 if i < 5 else -1) * abs(angle % (2 * 3.14159) - 3.14159) ** 0.5
                ))
            draw.polygon(points, fill=255)

        # 应用遮罩
        return Image.composite(image, Image.new('RGBA', image.size, (0, 0, 0, 0)), mask)

    @staticmethod
    def transform_image(image, x, y, width, height, angle):
        """应用缩放、旋转和位置变换"""
        import math
        
        # 确保宽度和高度大于0
        width = max(1, width)
        height = max(1, height)
        
        # 缩放
        scaled = image.resize((int(width), int(height)), Image.Resampling.LANCZOS)

        # 旋转
        if angle != 0:
            # 创建一个更大的画布以容纳旋转后的图像
            rotated_size = int(math.sqrt(width**2 + height**2)) * 2
            temp_canvas = Image.new('RGBA', (rotated_size, rotated_size), (0, 0, 0, 0))
            
            # 将缩放后的图像居中放置在临时画布上
            temp_x = (rotated_size - scaled.width) // 2
            temp_y = (rotated_size - scaled.height) // 2
            temp_canvas.paste(scaled, (temp_x, temp_y))
            
            # 旋转
            rotated = temp_canvas.rotate(angle, expand=True, resample=Image.Resampling.BILINEAR)
            
            # 计算新的偏移量
            offset_x = x - (rotated_size - scaled.width) // 2
            offset_y = y - (rotated_size - scaled.height) // 2
            
            return rotated, (offset_x, offset_y)
        else:
            # 如果没有旋转，直接返回缩放后的图像和位置
            return scaled, (x, y)

    @staticmethod
    def perspective_transform(image, vertices, output_size):
        """应用透视变换，将图像映射到指定的四边形顶点"""
        try:
            import numpy as np
            import cv2
            
            # 获取图像尺寸
            img_width, img_height = image.size
            
            # 源坐标：图像的四个角
            src_points = np.float32([
                [0, 0],
                [img_width, 0],
                [0, img_height],
                [img_width, img_height]
            ])
            
            # 目标坐标：指定的四个顶点
            # 确保顶点坐标是列表而不是元组
            dst_points = np.float32([list(v) for v in vertices])
            
            # 计算透视变换矩阵
            matrix = cv2.getPerspectiveTransform(src_points, dst_points)
            
            # 计算输出图像的尺寸
            output_width, output_height = output_size
            
            # 应用透视变换
            result = cv2.warpPerspective(
                np.array(image), 
                matrix, 
                (output_width, output_height),
                borderMode=cv2.BORDER_CONSTANT,
                borderValue=(0, 0, 0, 0)  # 透明背景
            )
            
            # 转换回PIL图像
            return Image.fromarray(result)
        except ImportError:
            # 如果没有安装opencv，返回原始图像
            return image

    @staticmethod
    def transform_with_vertices(image, vertices):
        """根据顶点坐标对图像进行变形"""
        try:
            import numpy as np
            import cv2
            
            # 获取图像尺寸
            img_width, img_height = image.size
            
            # 源坐标：图像的四个角
            src_points = np.float32([
                [0, 0],
                [img_width, 0],
                [0, img_height],
                [img_width, img_height]
            ])
            
            # 目标坐标：指定的四个顶点
            dst_points = np.float32([list(v) for v in vertices])
            
            # 计算包围盒
            min_x = min(v[0] for v in dst_points)
            min_y = min(v[1] for v in dst_points)
            max_x = max(v[0] for v in dst_points)
            max_y = max(v[1] for v in dst_points)
            
            # 调整目标坐标为相对于包围盒的坐标
            adjusted_dst_points = dst_points - np.array([min_x, min_y])
            
            # 计算输出尺寸
            output_width = int(max_x - min_x)
            output_height = int(max_y - min_y)
            
            # 确保输出尺寸大于0
            output_width = max(1, output_width)
            output_height = max(1, output_height)
            
            # 计算透视变换矩阵
            matrix = cv2.getPerspectiveTransform(src_points, adjusted_dst_points)
            
            # 应用透视变换，使用透明背景
            result = cv2.warpPerspective(
                np.array(image), 
                matrix, 
                (output_width, output_height),
                borderMode=cv2.BORDER_CONSTANT,
                borderValue=(0, 0, 0, 0)  # 透明背景而不是黑色背景
            )
            
            # 转换回PIL图像
            if len(result.shape) == 3 and result.shape[2] == 4:  # 如果有alpha通道
                return Image.fromarray(result), (min_x, min_y)
            else:  # 如果没有alpha通道
                return Image.fromarray(result).convert('RGBA'), (min_x, min_y)
        except ImportError:
            # 如果没有安装opencv，使用简单的变换
            min_x = min(v[0] for v in vertices)
            min_y = min(v[1] for v in vertices)
            # 确保返回的图像具有RGBA模式
            if image.mode != 'RGBA':
                image = image.convert('RGBA')
            return image, (min_x, min_y)

    @staticmethod
    def save_image(image, file_path, format='png', quality=95, keep_exif=False):
        """保存图像到文件"""
        try:
            # 如果是JPEG格式，转换为RGB
            if format.lower() == 'jpeg':
                image = image.convert('RGB')

            # 保存图像
            image.save(
                file_path,
                format=format.upper(),
                quality=quality,
                exif=image.info.get('exif') if keep_exif else None
            )
            return True
        except Exception as e:
            print(f"保存图像失败: {e}")
            return False