344 lines
14 KiB
Python
344 lines
14 KiB
Python
import re
|
||
from konabot.plugins.fx_process.color_handle import ColorHandle
|
||
import numpy as np
|
||
from PIL import Image, ImageDraw
|
||
from typing import List, Tuple, Dict, Optional
|
||
|
||
class GradientGenerator:
|
||
"""渐变生成器类"""
|
||
|
||
def __init__(self):
|
||
self.has_numpy = hasattr(np, '__version__')
|
||
|
||
def parse_color_list(self, color_list_str: str) -> List[Dict]:
|
||
"""解析渐变颜色列表字符串
|
||
|
||
Args:
|
||
color_list_str: 格式如 '[rgb(255,0,0)|(0,0)+rgb(0,255,0)|(0,100)+rgb(0,0,255)|(50,100)]'
|
||
|
||
Returns:
|
||
list: 包含颜色和位置信息的字典列表
|
||
"""
|
||
color_nodes = []
|
||
color_list_str = color_list_str.strip('[]').strip()
|
||
matches = color_list_str.split('+')
|
||
|
||
for single_str in matches:
|
||
color_str = single_str.split('|')[0]
|
||
pos_str = single_str.split('|')[1] if '|' in single_str else '0,0'
|
||
|
||
color = ColorHandle.parse_color(color_str.strip())
|
||
|
||
try:
|
||
pos_str = pos_str.replace('(', '').replace(')', '')
|
||
x_str, y_str = pos_str.split(',')
|
||
x_percent = float(x_str.strip().replace('%', ''))
|
||
y_percent = float(y_str.strip().replace('%', ''))
|
||
x_percent = max(0, min(100, x_percent))
|
||
y_percent = max(0, min(100, y_percent))
|
||
except:
|
||
x_percent = 0
|
||
y_percent = 0
|
||
|
||
color_nodes.append({
|
||
'color': color,
|
||
'position': (x_percent / 100.0, y_percent / 100.0)
|
||
})
|
||
|
||
if not color_nodes:
|
||
color_nodes = [
|
||
{'color': (255, 0, 0), 'position': (0, 0)},
|
||
{'color': (0, 0, 255), 'position': (1, 1)}
|
||
]
|
||
|
||
return color_nodes
|
||
|
||
def create_gradient(self, width: int, height: int, color_nodes: List[Dict]) -> Image.Image:
|
||
"""创建渐变图像
|
||
|
||
Args:
|
||
width: 图像宽度
|
||
height: 图像高度
|
||
color_nodes: 颜色节点列表
|
||
|
||
Returns:
|
||
Image.Image: 渐变图像
|
||
"""
|
||
if len(color_nodes) == 1:
|
||
return Image.new('RGB', (width, height), color_nodes[0]['color'])
|
||
elif len(color_nodes) == 2:
|
||
return self._create_linear_gradient(width, height, color_nodes)
|
||
else:
|
||
return self._create_radial_gradient(width, height, color_nodes)
|
||
|
||
def _create_linear_gradient(self, width: int, height: int, color_nodes: List[Dict]) -> Image.Image:
|
||
"""创建线性渐变"""
|
||
color1 = color_nodes[0]['color']
|
||
color2 = color_nodes[1]['color']
|
||
pos1 = color_nodes[0]['position']
|
||
pos2 = color_nodes[1]['position']
|
||
|
||
if self.has_numpy:
|
||
return self._create_linear_gradient_numpy(width, height, color1, color2, pos1, pos2)
|
||
else:
|
||
return self._create_linear_gradient_pil(width, height, color1, color2, pos1, pos2)
|
||
|
||
def _create_linear_gradient_numpy(self, width: int, height: int,
|
||
color1: Tuple, color2: Tuple,
|
||
pos1: Tuple, pos2: Tuple) -> Image.Image:
|
||
"""使用numpy创建线性渐变"""
|
||
# 创建坐标网格
|
||
x = np.linspace(0, 1, width)
|
||
y = np.linspace(0, 1, height)
|
||
xx, yy = np.meshgrid(x, y)
|
||
|
||
# 计算渐变方向
|
||
dx = pos2[0] - pos1[0]
|
||
dy = pos2[1] - pos1[1]
|
||
length_sq = dx * dx + dy * dy
|
||
|
||
if length_sq > 0:
|
||
# 计算投影参数
|
||
t = ((xx - pos1[0]) * dx + (yy - pos1[1]) * dy) / length_sq
|
||
t = np.clip(t, 0, 1)
|
||
else:
|
||
t = np.zeros_like(xx)
|
||
|
||
# 插值颜色
|
||
r = color1[0] + (color2[0] - color1[0]) * t
|
||
g = color1[1] + (color2[1] - color1[1]) * t
|
||
b = color1[2] + (color2[2] - color1[2]) * t
|
||
|
||
# 创建图像
|
||
gradient_array = np.stack([r, g, b], axis=-1).astype(np.uint8)
|
||
return Image.fromarray(gradient_array)
|
||
|
||
def _create_linear_gradient_pil(self, width: int, height: int,
|
||
color1: Tuple, color2: Tuple,
|
||
pos1: Tuple, pos2: Tuple) -> Image.Image:
|
||
"""使用PIL创建线性渐变(没有numpy时使用)"""
|
||
gradient = Image.new('RGB', (width, height))
|
||
draw = ImageDraw.Draw(gradient)
|
||
|
||
# 判断渐变方向
|
||
if abs(pos1[0] - pos2[0]) < 0.01: # 垂直渐变
|
||
y1 = int(pos1[1] * (height - 1))
|
||
y2 = int(pos2[1] * (height - 1))
|
||
|
||
if y2 < y1:
|
||
y1, y2 = y2, y1
|
||
color1, color2 = color2, color1
|
||
|
||
if y2 > y1:
|
||
for y in range(height):
|
||
if y <= y1:
|
||
fill_color = color1
|
||
elif y >= y2:
|
||
fill_color = color2
|
||
else:
|
||
ratio = (y - y1) / (y2 - y1)
|
||
r = int(color1[0] + (color2[0] - color1[0]) * ratio)
|
||
g = int(color1[1] + (color2[1] - color1[1]) * ratio)
|
||
b = int(color1[2] + (color2[2] - color1[2]) * ratio)
|
||
fill_color = (r, g, b)
|
||
|
||
draw.line([(0, y), (width, y)], fill=fill_color)
|
||
else:
|
||
draw.rectangle([0, 0, width, height], fill=color1)
|
||
|
||
elif abs(pos1[1] - pos2[1]) < 0.01: # 水平渐变
|
||
x1 = int(pos1[0] * (width - 1))
|
||
x2 = int(pos2[0] * (width - 1))
|
||
|
||
if x2 < x1:
|
||
x1, x2 = x2, x1
|
||
color1, color2 = color2, color1
|
||
|
||
if x2 > x1:
|
||
for x in range(width):
|
||
if x <= x1:
|
||
fill_color = color1
|
||
elif x >= x2:
|
||
fill_color = color2
|
||
else:
|
||
ratio = (x - x1) / (x2 - x1)
|
||
r = int(color1[0] + (color2[0] - color1[0]) * ratio)
|
||
g = int(color1[1] + (color2[1] - color1[1]) * ratio)
|
||
b = int(color1[2] + (color2[2] - color1[2]) * ratio)
|
||
fill_color = (r, g, b)
|
||
|
||
draw.line([(x, 0), (x, height)], fill=fill_color)
|
||
else:
|
||
draw.rectangle([0, 0, width, height], fill=color1)
|
||
|
||
else: # 对角渐变(简化处理为左上到右下)
|
||
for y in range(height):
|
||
for x in range(width):
|
||
distance = (x/width + y/height) / 2
|
||
r = int(color1[0] + (color2[0] - color1[0]) * distance)
|
||
g = int(color1[1] + (color2[1] - color1[1]) * distance)
|
||
b = int(color1[2] + (color2[2] - color1[2]) * distance)
|
||
draw.point((x, y), fill=(r, g, b))
|
||
|
||
return gradient
|
||
|
||
def _create_radial_gradient(self, width: int, height: int, color_nodes: List[Dict]) -> Image.Image:
|
||
"""创建径向渐变"""
|
||
if self.has_numpy and len(color_nodes) > 2:
|
||
return self._create_radial_gradient_numpy(width, height, color_nodes)
|
||
else:
|
||
return self._create_simple_gradient(width, height, color_nodes)
|
||
|
||
def _create_radial_gradient_numpy(self, width: int, height: int, color_nodes: List[Dict]) -> Image.Image:
|
||
"""使用numpy创建径向渐变(多色)"""
|
||
# 创建坐标网格
|
||
x = np.linspace(0, 1, width)
|
||
y = np.linspace(0, 1, height)
|
||
xx, yy = np.meshgrid(x, y)
|
||
|
||
# 提取颜色和位置
|
||
positions = np.array([node['position'] for node in color_nodes])
|
||
colors = np.array([node['color'] for node in color_nodes])
|
||
|
||
# 计算每个点到所有节点的距离
|
||
distances = np.sqrt((xx[:, :, np.newaxis] - positions[np.newaxis, np.newaxis, :, 0]) ** 2 +
|
||
(yy[:, :, np.newaxis] - positions[np.newaxis, np.newaxis, :, 1]) ** 2)
|
||
|
||
# 找到最近的两个节点
|
||
sorted_indices = np.argsort(distances, axis=2)
|
||
nearest_idx = sorted_indices[:, :, 0]
|
||
second_idx = sorted_indices[:, :, 1]
|
||
|
||
# 获取对应的颜色
|
||
nearest_colors = colors[nearest_idx]
|
||
second_colors = colors[second_idx]
|
||
|
||
# 获取距离并计算权重
|
||
nearest_dist = np.take_along_axis(distances, np.expand_dims(nearest_idx, axis=2), axis=2)[:, :, 0]
|
||
second_dist = np.take_along_axis(distances, np.expand_dims(second_idx, axis=2), axis=2)[:, :, 0]
|
||
|
||
total_dist = nearest_dist + second_dist
|
||
mask = total_dist > 0
|
||
weight1 = np.zeros_like(nearest_dist)
|
||
weight1[mask] = second_dist[mask] / total_dist[mask]
|
||
weight2 = 1 - weight1
|
||
|
||
# 插值颜色
|
||
r = nearest_colors[:, :, 0] * weight1 + second_colors[:, :, 0] * weight2
|
||
g = nearest_colors[:, :, 1] * weight1 + second_colors[:, :, 1] * weight2
|
||
b = nearest_colors[:, :, 2] * weight1 + second_colors[:, :, 2] * weight2
|
||
|
||
gradient_array = np.stack([r, g, b], axis=-1).astype(np.uint8)
|
||
return Image.fromarray(gradient_array)
|
||
|
||
def _create_simple_gradient(self, width: int, height: int, color_nodes: List[Dict]) -> Image.Image:
|
||
"""创建简化渐变(没有numpy或多色时使用)"""
|
||
gradient = Image.new('RGB', (width, height))
|
||
draw = ImageDraw.Draw(gradient)
|
||
|
||
if len(color_nodes) >= 2:
|
||
# 使用第一个和最后一个颜色创建简单渐变
|
||
color1 = color_nodes[0]['color']
|
||
color2 = color_nodes[-1]['color']
|
||
|
||
# 判断节点分布
|
||
x_positions = [node['position'][0] for node in color_nodes]
|
||
y_positions = [node['position'][1] for node in color_nodes]
|
||
|
||
if all(abs(x - x_positions[0]) < 0.01 for x in x_positions):
|
||
# 垂直渐变
|
||
for y in range(height):
|
||
ratio = y / (height - 1) if height > 1 else 0
|
||
r = int(color1[0] + (color2[0] - color1[0]) * ratio)
|
||
g = int(color1[1] + (color2[1] - color1[1]) * ratio)
|
||
b = int(color1[2] + (color2[2] - color1[2]) * ratio)
|
||
draw.line([(0, y), (width, y)], fill=(r, g, b))
|
||
else:
|
||
# 水平渐变
|
||
for x in range(width):
|
||
ratio = x / (width - 1) if width > 1 else 0
|
||
r = int(color1[0] + (color2[0] - color1[0]) * ratio)
|
||
g = int(color1[1] + (color2[1] - color1[1]) * ratio)
|
||
b = int(color1[2] + (color2[2] - color1[2]) * ratio)
|
||
draw.line([(x, 0), (x, height)], fill=(r, g, b))
|
||
else:
|
||
# 单色
|
||
draw.rectangle([0, 0, width, height], fill=color_nodes[0]['color'])
|
||
|
||
return gradient
|
||
|
||
def create_simple_gradient(self, width: int, height: int,
|
||
start_color: Tuple, end_color: Tuple,
|
||
direction: str = 'vertical') -> Image.Image:
|
||
"""创建简单双色渐变
|
||
|
||
Args:
|
||
width: 图像宽度
|
||
height: 图像高度
|
||
start_color: 起始颜色
|
||
end_color: 结束颜色
|
||
direction: 渐变方向 'vertical', 'horizontal', 'diagonal'
|
||
|
||
Returns:
|
||
Image.Image: 渐变图像
|
||
"""
|
||
if direction == 'vertical':
|
||
return self._create_vertical_gradient(width, height, start_color, end_color)
|
||
elif direction == 'horizontal':
|
||
return self._create_horizontal_gradient(width, height, start_color, end_color)
|
||
else: # diagonal
|
||
return self._create_diagonal_gradient(width, height, start_color, end_color)
|
||
|
||
def _create_vertical_gradient(self, width: int, height: int,
|
||
color1: Tuple, color2: Tuple) -> Image.Image:
|
||
"""创建垂直渐变"""
|
||
gradient = Image.new('RGB', (width, height))
|
||
draw = ImageDraw.Draw(gradient)
|
||
|
||
for y in range(height):
|
||
ratio = y / (height - 1) if height > 1 else 0
|
||
r = int(color1[0] + (color2[0] - color1[0]) * ratio)
|
||
g = int(color1[1] + (color2[1] - color1[1]) * ratio)
|
||
b = int(color1[2] + (color2[2] - color1[2]) * ratio)
|
||
draw.line([(0, y), (width, y)], fill=(r, g, b))
|
||
|
||
return gradient
|
||
|
||
def _create_horizontal_gradient(self, width: int, height: int,
|
||
color1: Tuple, color2: Tuple) -> Image.Image:
|
||
"""创建水平渐变"""
|
||
gradient = Image.new('RGB', (width, height))
|
||
draw = ImageDraw.Draw(gradient)
|
||
|
||
for x in range(width):
|
||
ratio = x / (width - 1) if width > 1 else 0
|
||
r = int(color1[0] + (color2[0] - color1[0]) * ratio)
|
||
g = int(color1[1] + (color2[1] - color1[1]) * ratio)
|
||
b = int(color1[2] + (color2[2] - color1[2]) * ratio)
|
||
draw.line([(x, 0), (x, height)], fill=(r, g, b))
|
||
|
||
return gradient
|
||
|
||
def _create_diagonal_gradient(self, width: int, height: int,
|
||
color1: Tuple, color2: Tuple) -> Image.Image:
|
||
"""创建对角渐变"""
|
||
if self.has_numpy:
|
||
return self._create_diagonal_gradient_numpy(width, height, color1, color2)
|
||
else:
|
||
return self._create_horizontal_gradient(width, height, color1, color2) # 降级为水平渐变
|
||
|
||
def _create_diagonal_gradient_numpy(self, width: int, height: int,
|
||
color1: Tuple, color2: Tuple) -> Image.Image:
|
||
"""使用numpy创建对角渐变"""
|
||
x = np.linspace(0, 1, width)
|
||
y = np.linspace(0, 1, height)
|
||
xx, yy = np.meshgrid(x, y)
|
||
|
||
distance = (xx + yy) / 2.0
|
||
|
||
r = color1[0] + (color2[0] - color1[0]) * distance
|
||
g = color1[1] + (color2[1] - color1[1]) * distance
|
||
b = color1[2] + (color2[2] - color1[2]) * distance
|
||
|
||
gradient_array = np.stack([r, g, b], axis=-1).astype(np.uint8)
|
||
return Image.fromarray(gradient_array) |