175 lines
6.3 KiB
Python
175 lines
6.3 KiB
Python
import vtracer
|
|
import tempfile
|
|
from wand.image import Image as WandImage
|
|
from PIL import Image as PILImage
|
|
import torch
|
|
import numpy as np
|
|
import os
|
|
|
|
class Vector:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"colormode": (["color", "binary"], {
|
|
"default": "color"
|
|
}),
|
|
"hierarchical": (["stacked", "cutout"], {
|
|
"default": "stacked"
|
|
}),
|
|
"mode": (["spline", "polygon", "none"], {
|
|
"default": "spline"
|
|
}),
|
|
"filter_speckle": ("INT", {
|
|
"default": 4,
|
|
"min": 1,
|
|
"max": 128,
|
|
"step": 1
|
|
}),
|
|
"color_precision": ("INT", {
|
|
"default": 6,
|
|
"min": 1,
|
|
"max": 10,
|
|
"step": 1
|
|
}),
|
|
"layer_difference": ("INT", {
|
|
"default": 16,
|
|
"min": 1,
|
|
"max": 20,
|
|
"step": 1
|
|
}),
|
|
"corner_threshold": ("INT", {
|
|
"default": 60,
|
|
"min": 1,
|
|
"max": 100,
|
|
"step": 1
|
|
}),
|
|
"length_threshold": ("FLOAT", {
|
|
"default": 4.0,
|
|
"min": 3.5,
|
|
"max": 10.0,
|
|
"step": 0.1
|
|
}),
|
|
"max_iterations": ("INT", {
|
|
"default": 10,
|
|
"min": 1,
|
|
"max": 20,
|
|
"step": 1
|
|
}),
|
|
"splice_threshold": ("INT", {
|
|
"default": 45,
|
|
"min": 1,
|
|
"max": 100,
|
|
"step": 1
|
|
}),
|
|
"path_precision": ("INT", {
|
|
"default": 8,
|
|
"min": 1,
|
|
"max": 10,
|
|
"step": 1
|
|
}),
|
|
"output_directory": ("STRING", {
|
|
"default": ""
|
|
})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("vector_image",)
|
|
|
|
FUNCTION = "process_image"
|
|
|
|
#OUTPUT_NODE = False
|
|
|
|
CATEGORY = "Image Vector"
|
|
|
|
|
|
|
|
def process_image(self, image, output_directory, colormode, hierarchical, mode, filter_speckle, color_precision,
|
|
layer_difference, corner_threshold, length_threshold, max_iterations, splice_threshold, path_precision):
|
|
try:
|
|
print(f"输入图像形状：{image.shape}")
|
|
|
|
# 确保输入 Tensor 是四维，具有 (B, C, H, W) 结构
|
|
if len(image.shape) == 4:
|
|
# 将图像从 (1, H, W, C) 转换为 (H, W, C)
|
|
image = image.squeeze(0)
|
|
|
|
print(f"处理后的图像形状：{image.shape}")
|
|
|
|
# 将 Tensor 转换为 PIL 图像
|
|
image_array = image.numpy()
|
|
image_array = (image_array * 255).astype(np.uint8)
|
|
pil_image = PILImage.fromarray(image_array)
|
|
|
|
print(f"转换后的 PIL 图像：{pil_image}")
|
|
|
|
# 将输入图像保存到临时文件
|
|
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_img_file:
|
|
pil_image.save(temp_img_file, format="PNG")
|
|
temp_img_path = temp_img_file.name
|
|
|
|
print(f"临时 PNG 文件路径：{temp_img_path}")
|
|
base_filename = os.path.splitext(os.path.basename(temp_img_path))[0]
|
|
|
|
# 检查并创建输出目录
|
|
if not output_directory:
|
|
output_directory = tempfile.gettempdir()
|
|
elif not os.path.exists(output_directory):
|
|
os.makedirs(output_directory)
|
|
|
|
# 生成 SVG 文件路径
|
|
svg_filename = os.path.join(output_directory, f"{base_filename}.svg")
|
|
|
|
print(f"SVG 文件路径：{svg_filename}")
|
|
|
|
# 使用 vtracer 将图像转换为 SVG
|
|
vtracer.convert_image_to_svg_py(temp_img_path,
|
|
svg_filename,
|
|
colormode=colormode,
|
|
hierarchical=hierarchical,
|
|
mode=mode,
|
|
filter_speckle=filter_speckle,
|
|
color_precision=color_precision,
|
|
layer_difference=layer_difference,
|
|
corner_threshold=corner_threshold,
|
|
length_threshold=length_threshold,
|
|
max_iterations=max_iterations,
|
|
splice_threshold=splice_threshold,
|
|
path_precision=path_precision)
|
|
|
|
print("SVG 文件已成功生成")
|
|
|
|
# 使用 wand 将 SVG 转换为 PNG
|
|
with WandImage(filename=svg_filename) as wand_image:
|
|
wand_image.format = 'png'
|
|
output_png_filename = os.path.join(output_directory, f"{base_filename}.png")
|
|
wand_image.save(filename=output_png_filename)
|
|
|
|
# 读取转换后的 PNG 图像
|
|
output_image = PILImage.open(output_png_filename).convert('RGB')
|
|
print(f"转换后的图像信息: 模式={output_image.mode}, 尺寸={output_image.size}")
|
|
|
|
# 将 PIL 图像转换为 torch tensor
|
|
tensor_image = torch.from_numpy(np.array(output_image)).permute(2, 0, 1).float() / 255
|
|
print(f"tensor 图像形状: {tensor_image.shape}, 最大值: {tensor_image.max()}, 最小值: {tensor_image.min()}")
|
|
|
|
|
|
return (tensor_image,)
|
|
|
|
except Exception as e:
|
|
print(f"错误：{str(e)}")
|
|
return (None,)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Vector": Vector
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Vector": "Vector"
|
|
}
|