Spaces:
Sleeping
Sleeping
File size: 7,420 Bytes
2cf467c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 | """
元素加载器
从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)
|