modules/infographics_generator/layout_system/element_loader.py

"""
元素加载器

从RGBA格式的PNG图片加载元素，根据图片尺寸确定元素尺寸，
并根据alpha通道生成mask
"""

import os
import numpy as np
from PIL import Image
from typing import Optional, Tuple

from .utils.nodes import LeafNode, NodeType


def resize_image_with_aspect_ratio(img: Image.Image, 
                                   target_width: int, 
                                   target_height: int) -> Tuple[Image.Image, Tuple[int, int]]:
    """
    保持横纵比resize图片，最短边对齐，不添加透明padding
    
    例如：原始1024x1024，目标500x1000，结果500x500（保持1:1比例，最短边对齐）
    
    Args:
        img: PIL图片对象（RGBA格式）
        target_width: 目标宽度
        target_height: 目标高度
    
    Returns:
        (resized_image, actual_size): 调整后的图片和实际尺寸
    """
    original_width, original_height = img.size
    original_aspect = original_width / original_height
    print("original_aspect: ", original_aspect)
    print("target_width: ", target_width, "target_height: ", target_height)
    # 找到目标尺寸的较短边，以较短边为基准
    if target_width <= target_height:
        # 目标宽度是较短边，以宽度为准
        new_width = target_width
        new_height = int(original_height * (target_width / original_width))
    else:
        # 目标高度是较短边，以高度为准
        new_height = target_height
        new_width = int(original_width * (target_height / original_height))
    
    # 确保不超过目标尺寸（双重检查）
    if new_width > target_width:
        new_width = target_width
        new_height = int(original_height * (target_width / original_width))
    if new_height > target_height:
        new_height = target_height
        new_width = int(original_width * (target_height / original_height))
    
    print("new_width: ", new_width, "new_height: ", new_height)
    # Resize图片（保持横纵比，不添加透明padding）
    resized_img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
    
    return resized_img, (new_width, new_height)


class ElementLoader:
    """元素加载器"""
    
    def __init__(self, base_dir: Optional[str] = None):
        """
        初始化元素加载器
        
        Args:
            base_dir: 图片文件的基础目录，如果提供，相对路径会基于此目录
        """
        self.base_dir = base_dir
    
    def load_from_image(self, image_path: str, node_id: str, 
                       node_type: NodeType, 
                       metadata: Optional[dict] = None,
                       target_width: Optional[float] = None,
                       target_height: Optional[float] = None) -> LeafNode:
        """
        从PNG图片加载元素
        
        Args:
            image_path: 图片文件路径（可以是绝对路径或相对于base_dir的路径）
            node_id: 节点ID
            node_type: 节点类型
            metadata: 额外的元数据
            target_width: 目标宽度（可选，如果提供则resize图片）
            target_height: 目标高度（可选，如果提供则resize图片）
            
        Returns:
            LeafNode对象，包含图片尺寸和mask
        """
        # 解析路径
        full_path = self._resolve_path(image_path)
        
        # 加载图片
        img = Image.open(full_path)
        
        # 确保是RGBA格式
        if img.mode != 'RGBA':
            img = img.convert('RGBA')
        
        # 获取原始尺寸
        original_width, original_height = img.size
        
        # 如果指定了目标尺寸，resize图片（保持横纵比，最短边对齐）
        if target_width is not None and target_height is not None:
            target_width = int(target_width)
            target_height = int(target_height)
            # 使用保持横纵比的resize函数（最短边对齐，不添加透明padding）
            img, (width, height) = resize_image_with_aspect_ratio(
                img, target_width, target_height
            )
        else:
            width, height = original_width, original_height
        
        # 从alpha通道生成mask
        # mask是二值图像，alpha > 0 的像素为1，否则为0
        alpha_channel = np.array(img.split()[3])  # 获取alpha通道
        mask = (alpha_channel > 0).astype(np.uint8) * 255
        
        # 创建元数据
        node_metadata = metadata or {}
        node_metadata['image_path'] = image_path
        node_metadata['image_size'] = (width, height)
        node_metadata['original_image_size'] = (original_width, original_height)
        if target_width is not None and target_height is not None:
            node_metadata['resized'] = True
        
        # 创建叶子节点
        node = LeafNode(
            node_id=node_id,
            node_type=node_type,
            width=float(width),
            height=float(height),
            mask=mask,
            metadata=node_metadata
        )
        
        return node
    
    def _resolve_path(self, image_path: str) -> str:
        """解析图片路径"""
        if os.path.isabs(image_path):
            return image_path
        
        if self.base_dir:
            return os.path.join(self.base_dir, image_path)
        
        return image_path
    
    def load_chart(self, image_path: str, node_id: str, 
                   metadata: Optional[dict] = None) -> LeafNode:
        """加载图表元素"""
        return self.load_from_image(image_path, node_id, NodeType.CHART, metadata)
    
    def load_image(self, image_path: str, node_id: str,
                  metadata: Optional[dict] = None) -> LeafNode:
        """加载图像元素"""
        return self.load_from_image(image_path, node_id, NodeType.IMAGE, metadata)
    
    def load_text(self, image_path: str, node_id: str,
                 metadata: Optional[dict] = None) -> LeafNode:
        """加载文本元素（文本渲染为图片）"""
        return self.load_from_image(image_path, node_id, NodeType.TEXT, metadata)
    
    def load_shape(self, image_path: str, node_id: str,
                  metadata: Optional[dict] = None) -> LeafNode:
        """加载形状元素"""
        return self.load_from_image(image_path, node_id, NodeType.SHAPE, metadata)


def load_element_from_image(image_path: str, node_id: str, 
                           node_type: NodeType,
                           base_dir: Optional[str] = None,
                           metadata: Optional[dict] = None,
                           target_width: Optional[float] = None,
                           target_height: Optional[float] = None) -> LeafNode:
    """
    便捷函数：从图片加载元素
    
    Args:
        image_path: 图片文件路径
        node_id: 节点ID
        node_type: 节点类型
        base_dir: 基础目录（可选）
        metadata: 额外的元数据（可选）
        target_width: 目标宽度（可选，如果提供则resize图片）
        target_height: 目标高度（可选，如果提供则resize图片）
        
    Returns:
        LeafNode对象
    """
    loader = ElementLoader(base_dir=base_dir)
    return loader.load_from_image(image_path, node_id, node_type, metadata, 
                                 target_width, target_height)