nodes/rickrender-ComfyUI-Vectorizer-API
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update vectorizer_ai.py

Browse Source 
Added an additional node to remove the single largest shape from a returned result, as well as some additional input/output options for saving.
This commit is contained in:
renderartist
2025-08-19 14:41:11 -07:00
committed by GitHub
parent 624f7626bc
commit df1d17009d
1 changed files with 146 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vectorizer_ai.py
+146 -53
View File
@@ -7,12 +7,18 @@ import json
from PIL import Image
import folder_paths
try:
from lxml import etree
except ImportError:
print("Warning: lxml is not installed. The 'Background Remover (Shape)' node will not work.")
print("Please install it by running: pip install lxml")
etree = None
try:
import cairosvg
except ImportError:
print("Warning: cairosvg is not installed. The 'scaled_png' output format will not work.")
print("Please install it by running: pip install cairosvg")
print("On Debian/Ubuntu, you may also need to run: sudo apt-get install libcairo2-dev pkg-config")
cairosvg = None
class VectorizerAINode:
@@ -47,18 +53,21 @@ class VectorizerAINode:
"scale": ("FLOAT", {"default": 4.0, "min": 1.0, "max": 16.0, "step": 0.5}),
},
"optional": {
"filename_prefix": ("STRING", {"default": "vectorized/vector"}),
"save_svg": ("BOOLEAN", {"default": True}),
"filename_prefix": ("STRING", {"default": "SVG/vector"}),
"max_colors": ("INT", {"default": 0, "min": 0, "max": 256, "step": 1, "display": "slider"}),
"min_shape_area": ("FLOAT", {"default": 0.125, "min": 0.0, "max": 100.0, "step": 0.001}),
"adobe_compatibility": ("BOOLEAN", {"default": False}),
"disable_gap_filler": ("BOOLEAN", {"default": True}),
}
}
RETURN_TYPES = ("IMAGE",)
RETURN_TYPES = ("IMAGE", "STRING",)
RETURN_NAMES = ("image", "svg_data",)
FUNCTION = "process_vectorization"
CATEGORY = "Conversion"
def process_vectorization(self, image, api_id, api_secret, output_format, mode, scale, filename_prefix="vectorized/vector", max_colors=0, min_shape_area=0.125, adobe_compatibility=False):
def process_vectorization(self, image, api_id, api_secret, output_format, mode, scale, save_svg=True, filename_prefix="vectorized/vector", max_colors=0, min_shape_area=0.125, adobe_compatibility=False, disable_gap_filler=True):
final_api_id = api_id
final_api_secret = api_secret
@@ -73,20 +82,22 @@ class VectorizerAINode:
if not final_api_id or final_api_id == "YOUR_API_ID_HERE":
print("Vectorizer.AI Error: API ID is missing. Please provide it in the UI or in config.json.")
return (image,)
return (image, "",)
pil_image = Image.fromarray((255. * image[0].cpu().numpy()).astype(np.uint8))
img_buffer = io.BytesIO()
pil_image.save(img_buffer, format='PNG')
img_buffer.seek(0)
api_output_format = 'svg' if output_format == 'scaled_png' else output_format
api_output_format = 'png' if output_format == 'png' else 'svg'
data = {"mode": mode, "output.file_format": api_output_format}
if max_colors > 0: data["processing.max_colors"] = max_colors
if min_shape_area > 0: data["processing.shapes.min_area_px"] = min_shape_area
if api_output_format == 'svg' and adobe_compatibility: data["output.svg.adobe_compatibility_mode"] = "true"
if adobe_compatibility: data["output.svg.adobe_compatibility_mode"] = "true"
if disable_gap_filler: data["output.gap_filler.enabled"] = "false"
print("Sending image to Vectorizer.AI for processing...")
print(f"Requesting '{api_output_format}' from Vectorizer.AI...")
try:
response = requests.post(
'https://vectorizer.ai/api/v1/vectorize',
@@ -97,45 +108,46 @@ class VectorizerAINode:
response.raise_for_status()
except requests.exceptions.RequestException as e:
print(f"Vectorizer.AI API Error: {e}")
return (image,)
return (image, "",)
print("Processing complete.")
full_output_folder, filename, counter, subfolder, _ = folder_paths.get_save_image_path(filename_prefix, self.output_dir)
if output_format == 'scaled_png':
if cairosvg is None: return (image,)
svg_filename = f"{filename}_{counter:05}.svg"
svg_path = os.path.join(full_output_folder, svg_filename)
with open(svg_path, 'wb') as f: f.write(response.content)
print(f"Saved vector file to: {svg_path}")
try:
print(f"Scaling SVG to {scale}x PNG...")
png_bytes = cairosvg.svg2png(bytestring=response.content, scale=scale)
png_filename = f"{filename}_{counter:05}_scaled_{scale}x.png"
png_path = os.path.join(full_output_folder, png_filename)
with open(png_path, 'wb') as f: f.write(png_bytes)
print(f"Saved scaled PNG to: {png_path}")
processed_pil_image = Image.open(io.BytesIO(png_bytes)).convert("RGB")
processed_np_image = np.array(processed_pil_image).astype(np.float32) / 255.0
return (torch.from_numpy(processed_np_image)[None,],)
except Exception as e:
print(f"ERROR during SVG to PNG conversion: {e}")
return (image,)
image_out = image
svg_data_string = ""
elif output_format == 'png':
file = f"{filename}_{counter:05}.png"
path = os.path.join(full_output_folder, file)
with open(path, 'wb') as f: f.write(response.content)
print(f"Saved API-generated PNG to: {path}")
if api_output_format == 'svg':
svg_data_string = response.content.decode('utf-8')
if save_svg:
svg_path = os.path.join(full_output_folder, f"{filename}_{counter:05}.svg")
with open(svg_path, 'w', encoding='utf-8') as f:
f.write(svg_data_string)
print(f"Saved original SVG to: {svg_path}")
if output_format == 'scaled_png':
if cairosvg is None: return (image, svg_data_string)
try:
print(f"Scaling SVG to {scale}x PNG...")
png_bytes = cairosvg.svg2png(bytestring=response.content, scale=scale)
png_filename = f"{filename}_{counter:05}_scaled_{scale}x.png"
png_path = os.path.join(full_output_folder, png_filename)
with open(png_path, 'wb') as f: f.write(png_bytes)
print(f"Saved scaled PNG to: {png_path}")
processed_pil_image = Image.open(io.BytesIO(png_bytes)).convert("RGB")
processed_np_image = np.array(processed_pil_image).astype(np.float32) / 255.0
image_out = torch.from_numpy(processed_np_image)[None,]
except Exception as e:
print(f"ERROR during SVG to PNG conversion: {e}")
elif api_output_format == 'png':
png_path = os.path.join(full_output_folder, f"{filename}_{counter:05}.png")
with open(png_path, 'wb') as f: f.write(response.content)
print(f"Saved API-generated PNG to: {png_path}")
processed_pil_image = Image.open(io.BytesIO(response.content)).convert("RGB")
processed_np_image = np.array(processed_pil_image).astype(np.float32) / 255.0
return (torch.from_numpy(processed_np_image)[None,],)
else:
file = f"{filename}_{counter:05}.svg"
path = os.path.join(full_output_folder, file)
with open(path, 'wb') as f: f.write(response.content)
print(f"Saved vector file to: {path}")
return (image,)
image_out = torch.from_numpy(processed_np_image)[None,]
return (image_out, svg_data_string,)
class BackgroundRemoverNode:
@classmethod
@@ -147,7 +159,6 @@ class BackgroundRemoverNode:
"invert_mask": ("BOOLEAN", {"default": False}),
}
}
RETURN_TYPES = ("IMAGE", "MASK",)
RETURN_NAMES = ("image_rgba", "mask",)
FUNCTION = "remove_background"
@@ -155,39 +166,121 @@ class BackgroundRemoverNode:
def remove_background(self, image, threshold, invert_mask):
img_np = np.clip(255. * image[0].cpu().numpy(), 0, 255).astype(np.uint8)
top_border = img_np[0, :, :]
bottom_border = img_np[-1, :, :]
left_border = img_np[:, 0, :]
right_border = img_np[:, -1, :]
border_pixels = np.concatenate([top_border, bottom_border, left_border, right_border], axis=0)
key_color = np.mean(border_pixels, axis=0)
print(f"Automatically detected background color: R={int(key_color[0])}, G={int(key_color[1])}, B={int(key_color[2])}")
distances = np.linalg.norm(img_np.astype(np.float32) - key_color.astype(np.float32), axis=2)
threshold_value = threshold * 255.0
mask_bool = distances > threshold_value
if invert_mask:
mask_bool = ~mask_bool
if invert_mask: mask_bool = ~mask_bool
mask_float = mask_bool.astype(np.float32)
alpha_channel = (mask_float * 255).astype(np.uint8)
rgba_np = np.dstack((img_np, alpha_channel))
image_rgba_tensor = torch.from_numpy(rgba_np.astype(np.float32) / 255.0)[None,]
mask_tensor = torch.from_numpy(mask_float)[None,]
return (image_rgba_tensor, mask_tensor)
class BackgroundRemoverSVGNode:
def __init__(self):
self.output_dir = folder_paths.get_output_directory()
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"svg_data": ("STRING", {"forceInput": True}),
"scale": ("FLOAT", {"default": 4.0, "min": 1.0, "max": 16.0, "step": 0.5}),
},
"optional": {
"save_svg": ("BOOLEAN", {"default": False}),
"filename_prefix": ("STRING", {"default": "SVG/vector_edited"}),
}
}
RETURN_TYPES = ("IMAGE", "MASK", "STRING",)
RETURN_NAMES = ("image_rgba", "mask", "modified_svg_data",)
FUNCTION = "remove_background_shape"
CATEGORY = "Conversion"
def remove_background_shape(self, svg_data, scale, save_svg=False, filename_prefix="SVG/vector_edited"):
blank_image = torch.zeros((1, 128, 128, 4), dtype=torch.float32)
blank_mask = torch.zeros((1, 128, 128), dtype=torch.float32)
blank_svg = ""
if not svg_data or etree is None or cairosvg is None:
return (blank_image, blank_mask, blank_svg)
try:
parser = etree.XMLParser(recover=True)
root = etree.fromstring(svg_data.encode('utf-8'), parser=parser)
paths = root.xpath('//svg:path', namespaces={'svg': 'http://www.w3.org/2000/svg'})
largest_path = None
largest_area = -1
for path in paths:
d = path.get('d')
if not d: continue
try:
coords = [float(c) for c in d.replace('M', ' ').replace('L', ' ').replace('C', ' ').replace('Z', ' ').replace('z', ' ').replace(',', ' ').split() if c]
if len(coords) < 2: continue
x_coords, y_coords = coords[0::2], coords[1::2]
min_x, max_x = min(x_coords), max(x_coords)
min_y, max_y = min(y_coords), max(y_coords)
area = (max_x - min_x) * (max_y - min_y)
if area > largest_area:
largest_area = area
largest_path = path
except ValueError:
continue
if largest_path is not None:
print(f"Found and removed largest shape with area: {largest_area}")
largest_path.getparent().remove(largest_path)
else:
print("Warning: Could not identify a largest shape to remove.")
modified_svg_data = etree.tostring(root, pretty_print=True).decode('utf-8')
if save_svg:
full_output_folder, filename, counter, subfolder, _ = folder_paths.get_save_image_path(filename_prefix, self.output_dir)
file = f"{filename}_{counter:05}.svg"
path = os.path.join(full_output_folder, file)
with open(path, 'w', encoding='utf-8') as f:
f.write(modified_svg_data)
print(f"Saved cleaned SVG file to: {path}")
print(f"Rendering modified SVG to {scale}x PNG...")
png_bytes = cairosvg.svg2png(bytestring=modified_svg_data.encode('utf-8'), scale=scale)
pil_image = Image.open(io.BytesIO(png_bytes)).convert("RGBA")
rgba_np = np.array(pil_image).astype(np.float32) / 255.0
mask_float = rgba_np[..., 3]
image_rgba_tensor = torch.from_numpy(rgba_np)[None,]
mask_tensor = torch.from_numpy(mask_float)[None,]
return (image_rgba_tensor, mask_tensor, modified_svg_data)
except Exception as e:
print(f"ERROR during SVG background removal: {e}")
return (blank_image, blank_mask, blank_svg)
NODE_CLASS_MAPPINGS = {
"VectorizerAINode": VectorizerAINode,
"BackgroundRemoverNode": BackgroundRemoverNode
"BackgroundRemoverNode": BackgroundRemoverNode,
"BackgroundRemoverSVGNode": BackgroundRemoverSVGNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"VectorizerAINode": "Vectorizer.ai API",
"BackgroundRemoverNode": "Background Remover (Color)"
"BackgroundRemoverNode": "Background Remover (Color)",
"BackgroundRemoverSVGNode": "Background Remover (Shape)",
}