Files
konabot/konabot/plugins/fx_process/gradient.py
2025-12-09 20:21:36 +08:00

344 lines
14 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

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)