diff --git a/AIvector.py b/AIvector.py new file mode 100644 index 0000000..cc3498d --- /dev/null +++ b/AIvector.py @@ -0,0 +1,216 @@ +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 AIvector: + """ + A example node + + Class methods + ------------- + INPUT_TYPES (dict): + Tell the main program input parameters of nodes. + IS_CHANGED: + optional method to control when the node is re executed. + + Attributes + ---------- + RETURN_TYPES (`tuple`): + The type of each element in the output tulple. + RETURN_NAMES (`tuple`): + Optional: The name of each output in the output tulple. + FUNCTION (`str`): + The name of the entry-point method. For example, if `FUNCTION = "execute"` then it will run Example().execute() + OUTPUT_NODE ([`bool`]): + If this node is an output node that outputs a result/image from the graph. The SaveImage node is an example. + The backend iterates on these output nodes and tries to execute all their parents if their parent graph is properly connected. + Assumed to be False if not present. + CATEGORY (`str`): + The category the node should appear in the UI. + execute(s) -> tuple || None: + The entry point method. The name of this method must be the same as the value of property `FUNCTION`. + For example, if `FUNCTION = "execute"` then this method's name must be `execute`, if `FUNCTION = "foo"` then it must be `foo`. + """ + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + """ + Return a dictionary which contains config for all input fields. + Some types (string): "MODEL", "VAE", "CLIP", "CONDITIONING", "LATENT", "IMAGE", "INT", "STRING", "FLOAT". + Input types "INT", "STRING" or "FLOAT" are special values for fields on the node. + The type can be a list for selection. + + Returns: `dict`: + - Key input_fields_group (`string`): Can be either required, hidden or optional. A node class must have property `required` + - Value input_fields (`dict`): Contains input fields config: + * Key field_name (`string`): Name of a entry-point method's argument + * Value field_config (`tuple`): + + First value is a string indicate the type of field or a list for selection. + + Secound value is a config for type "INT", "STRING" or "FLOAT". + """ + 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": 10, + "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 = "Vectorization" + + + + 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 = { + "AIvector": AIvector +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "AIvector": "AIvector" +} \ No newline at end of file diff --git a/example/AIvector.json b/example/AIvector.json new file mode 100644 index 0000000..77e6c1e --- /dev/null +++ b/example/AIvector.json @@ -0,0 +1,150 @@ +{ + "last_node_id": 84, + "last_link_id": 162, + "nodes": [ + { + "id": 82, + "type": "LoadImage", + "pos": [ + 284.0366204412037, + 814.6605073257773 + ], + "size": [ + 315, + 314.00001525878906 + ], + "flags": {}, + "order": 0, + "mode": 0, + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 161 + ], + "shape": 3 + }, + { + "name": "MASK", + "type": "MASK", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "18035-633794744-jdsh, black_background, animal, a colorful animal, colorful painting, texture _lora_jdsh-000008_0.9_.png", + "image" + ] + }, + { + "id": 83, + "type": "AIvector", + "pos": [ + 715, + 808 + ], + "size": { + "0": 315, + "1": 322 + }, + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 161, + "slot_index": 0 + } + ], + "outputs": [ + { + "name": "vector_image", + "type": "IMAGE", + "links": [ + 162 + ], + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "AIvector" + }, + "widgets_values": [ + "color", + "stacked", + "spline", + 4, + 6, + 16, + 60, + 4, + 10, + 45, + 8, + "D:\\AARG\\AARG" + ] + }, + { + "id": 84, + "type": "PreviewImage", + "pos": [ + 1147, + 803 + ], + "size": [ + 210, + 246 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 162, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "PreviewImage" + } + } + ], + "links": [ + [ + 161, + 82, + 0, + 83, + 0, + "IMAGE" + ], + [ + 162, + 83, + 0, + 84, + 0, + "IMAGE" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.9229599817709203, + "offset": [ + -262.16629897186385, + -591.0167384245041 + ] + } + }, + "version": 0.4 +} \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..666c661 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,5 @@ +vtracer==0.3.2 +Wand==0.6.10 +Pillow==9.4.0 +torch==2.0.0 +numpy==1.24.2 \ No newline at end of file