diff --git a/README.md b/README.md new file mode 100644 index 0000000..83b59e2 --- /dev/null +++ b/README.md @@ -0,0 +1,57 @@ +# TJ_Vector - ComfyUI Custom Node + +ラインアート画像から中心線を抽出し、SVGパスとして出力するComfyUIカスタムノード + +## ノード + +### 1. Line Region Detector +ラインアート画像から線領域を検出 + +**入力:** +- `image`: 入力画像 +- `background_mode`: 背景検出モード (auto/white/black) +- `line_detection_method`: 線検出方法 (edge/morphology/hybrid) +- その他のパラメータ + +**出力:** +- `line_mask`: 線領域のマスク +- `fill_mask`: 塗り領域のマスク +- `preview_image`: プレビュー画像 +- `color_info`: 色情報 (JSON) + +### 2. Centerline to SVG +線マスクから中心線を抽出してSVGパスを生成 + +**入力:** +- `line_mask`: 線領域のマスク +- `algorithm`: 中心線抽出アルゴリズム (ridge/skeleton/medial_axis) +- `smoothing`: パス平滑化レベル +- その他のパラメータ + +**出力:** +- `svg_string`: SVG形式の文字列 +- `centerline_image`: 中心線プレビュー +- `statistics`: 統計情報 (JSON) + +## インストール + +1. このフォルダを `ComfyUI/custom_nodes/` にコピー +2. 依存関係をインストール: +```bash +pip install -r requirements.txt +``` + +## 使い方 + +1. ComfyUIを起動 +2. ノードメニューから `TJ_Vector` カテゴリを選択 +3. `Line Region Detector` ノードで線を検出 +4. `Centerline to SVG` ノードでSVGを生成 + +## ライセンス + +MIT License + +## バージョン + +v1.0 - 初期リリース diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..44843de --- /dev/null +++ b/__init__.py @@ -0,0 +1,20 @@ +""" +TJ_Vector - ComfyUI Custom Nodes +Line art to vector conversion nodes +""" + +from .line_region_detector import NODE_CLASS_MAPPINGS as LINE_NODES, NODE_DISPLAY_NAME_MAPPINGS as LINE_NAMES +from .centerline_to_svg import NODE_CLASS_MAPPINGS as SVG_NODES, NODE_DISPLAY_NAME_MAPPINGS as SVG_NAMES + +# Combine all node mappings +NODE_CLASS_MAPPINGS = { + **LINE_NODES, + **SVG_NODES +} + +NODE_DISPLAY_NAME_MAPPINGS = { + **LINE_NAMES, + **SVG_NAMES +} + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/centerline_to_svg.py b/centerline_to_svg.py new file mode 100644 index 0000000..f832b1f --- /dev/null +++ b/centerline_to_svg.py @@ -0,0 +1,398 @@ +""" +Centerline to SVG Node for ComfyUI +Extracts centerlines from line masks and generates SVG paths +""" + +import numpy as np +import torch +import cv2 +from scipy import ndimage +from scipy.interpolate import splprep, splev +from skimage.morphology import medial_axis, skeletonize +import json +from datetime import datetime + + +class CenterlineToSVG: + """ + Extracts centerlines from line masks and generates SVG paths + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "line_mask": ("MASK",), + "algorithm": (["ridge", "skeleton", "medial_axis"],), + "smoothing": ("FLOAT", { + "default": 2.0, + "min": 0.0, + "max": 10.0, + "step": 0.1 + }), + "min_path_length": ("INT", { + "default": 10, + "min": 1, + "max": 1000, + "step": 1 + }), + "simplify_tolerance": ("FLOAT", { + "default": 1.0, + "min": 0.0, + "max": 10.0, + "step": 0.1 + }), + "bezier_smoothing": ("BOOLEAN", {"default": True}), + }, + "optional": { + "original_image": ("IMAGE",), + "color_info": ("STRING", {"default": ""}), + "preserve_colors": ("BOOLEAN", {"default": True}), + } + } + + RETURN_TYPES = ("STRING", "IMAGE", "STRING") + RETURN_NAMES = ("svg_string", "centerline_image", "statistics") + FUNCTION = "generate_svg" + CATEGORY = "TJ_Vector" + OUTPUT_NODE = True + + def generate_svg(self, line_mask, algorithm="ridge", smoothing=2.0, + min_path_length=10, simplify_tolerance=1.0, bezier_smoothing=True, + original_image=None, color_info="", preserve_colors=True): + """ + Main function to generate SVG from line mask + """ + # Convert mask to numpy + mask_np = line_mask[0].cpu().numpy() + mask_255 = (mask_np * 255).astype(np.uint8) + + height, width = mask_255.shape + + # Step 1: Extract centerline + centerline = self._extract_centerline(mask_255, algorithm) + + # Step 2: Generate paths from centerline + paths = self._extract_paths(centerline, min_path_length) + + # Step 3: Simplify paths + simplified_paths = self._simplify_paths(paths, simplify_tolerance) + + # Step 4: Smooth paths + if bezier_smoothing: + smoothed_paths = self._smooth_paths(simplified_paths, smoothing) + else: + smoothed_paths = simplified_paths + + # Step 5: Extract colors (optional) + path_colors = self._get_path_colors( + smoothed_paths, original_image, color_info, preserve_colors + ) + + # Step 6: Generate SVG + svg_string = self._generate_svg_string( + smoothed_paths, path_colors, width, height + ) + + # Step 7: Create preview + preview_image = self._create_centerline_preview( + centerline, smoothed_paths, width, height + ) + preview_tensor = torch.from_numpy(preview_image.astype(np.float32) / 255.0).unsqueeze(0) + + # Step 8: Generate statistics + statistics = self._generate_statistics( + smoothed_paths, path_colors, width, height, algorithm + ) + statistics_str = json.dumps(statistics, indent=2) + + return (svg_string, preview_tensor, statistics_str) + + def _extract_centerline(self, mask, algorithm): + """ + Extract centerline using specified algorithm + """ + if np.sum(mask > 127) == 0: + return np.zeros_like(mask) + + if algorithm == "skeleton": + return self._skeleton_method(mask) + elif algorithm == "medial_axis": + return self._medial_axis_method(mask) + elif algorithm == "ridge": + return self._ridge_detection_method(mask) + + return mask + + def _skeleton_method(self, mask): + """ + Zhang-Suen thinning algorithm + """ + binary = mask > 127 + skeleton = skeletonize(binary) + return (skeleton * 255).astype(np.uint8) + + def _medial_axis_method(self, mask): + """ + Medial axis transform (more accurate for thick lines) + """ + binary = mask > 127 + skel, distance = medial_axis(binary, return_distance=True) + + # Filter by distance to remove noise + if np.any(distance > 0): + threshold = np.percentile(distance[distance > 0], 10) + skel_filtered = skel & (distance > threshold) + return (skel_filtered * 255).astype(np.uint8) + + return (skel * 255).astype(np.uint8) + + def _ridge_detection_method(self, mask): + """ + Ridge detection using distance transform + """ + binary = mask > 127 + + # Distance transform + dist_transform = cv2.distanceTransform(mask, cv2.DIST_L2, 5) + + # Normalize + if np.max(dist_transform) > 0: + dist_norm = dist_transform / np.max(dist_transform) + else: + dist_norm = dist_transform + + # Detect local maxima (ridges) + # Use morphological gradient to find ridges + kernel = np.ones((3, 3), np.uint8) + + # Dilate and erode to find local maxima + dilated = cv2.dilate(dist_norm, kernel, iterations=1) + ridge = (dist_norm == dilated) & (dist_norm > 0.1) + + # Thin the ridge + ridge_uint8 = (ridge * 255).astype(np.uint8) + thinned = cv2.ximgproc.thinning(ridge_uint8, thinningType=cv2.ximgproc.THINNING_ZHANGSUEN) + + return thinned + + def _extract_paths(self, centerline, min_length): + """ + Extract individual paths from centerline + """ + # Find contours + contours, _ = cv2.findContours(centerline, cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE) + + paths = [] + for contour in contours: + if len(contour) >= min_length: + # Convert to simple point list + points = contour.squeeze() + if len(points.shape) == 1: + points = points.reshape(1, -1) + paths.append(points) + + return paths + + def _simplify_paths(self, paths, tolerance): + """ + Simplify paths using Douglas-Peucker algorithm + """ + if tolerance <= 0: + return paths + + simplified = [] + for path in paths: + # Reshape for cv2.approxPolyDP + path_reshaped = path.reshape(-1, 1, 2).astype(np.float32) + approx = cv2.approxPolyDP(path_reshaped, tolerance, closed=False) + simplified_path = approx.squeeze() + + if len(simplified_path.shape) == 1: + simplified_path = simplified_path.reshape(1, -1) + + if len(simplified_path) >= 2: + simplified.append(simplified_path) + + return simplified + + def _smooth_paths(self, paths, smoothing): + """ + Smooth paths using spline interpolation + """ + if smoothing <= 0: + return paths + + smoothed = [] + for path in paths: + if len(path) < 4: + smoothed.append(path) + continue + + try: + # Parametric spline interpolation + x = path[:, 0] + y = path[:, 1] + + # Remove duplicate points + unique_indices = np.where(np.sum(np.abs(np.diff(path, axis=0)), axis=1) > 0)[0] + if len(unique_indices) < 3: + smoothed.append(path) + continue + + unique_indices = np.concatenate([[0], unique_indices + 1]) + x_unique = x[unique_indices] + y_unique = y[unique_indices] + + if len(x_unique) < 4: + smoothed.append(path) + continue + + # Fit spline + tck, u = splprep([x_unique, y_unique], s=smoothing * len(x_unique), k=min(3, len(x_unique) - 1)) + + # Generate smooth curve + u_new = np.linspace(0, 1, len(path)) + x_new, y_new = splev(u_new, tck) + + smooth_path = np.column_stack([x_new, y_new]) + smoothed.append(smooth_path) + except Exception as e: + # If smoothing fails, use original path + smoothed.append(path) + + return smoothed + + def _get_path_colors(self, paths, original_image, color_info, preserve_colors): + """ + Get colors for each path + """ + if not preserve_colors or original_image is None: + return ["#000000"] * len(paths) + + # Convert image to numpy + img_np = original_image[0].cpu().numpy() + img_255 = (img_np * 255).astype(np.uint8) + + colors = [] + for path in paths: + # Sample color from middle of path + mid_idx = len(path) // 2 + x, y = path[mid_idx].astype(int) + + # Clamp coordinates + y = max(0, min(y, img_255.shape[0] - 1)) + x = max(0, min(x, img_255.shape[1] - 1)) + + # Get color + if len(img_255.shape) >= 3: + r, g, b = img_255[y, x, :3] + else: + r = g = b = img_255[y, x] + + hex_color = "#{:02x}{:02x}{:02x}".format(int(r), int(g), int(b)) + colors.append(hex_color) + + return colors + + def _generate_svg_string(self, paths, colors, width, height): + """ + Generate SVG string from paths + """ + timestamp = datetime.now().isoformat() + + svg_lines = [ + f'', + ' ', + f' {timestamp}', + ' tj_comfyui_centerVector v1.0', + ' ', + ' ', + ] + + for i, (path, color) in enumerate(zip(paths, colors)): + svg_path = self._path_to_svg_d(path) + svg_lines.append( + f' ' + ) + + svg_lines.append(' ') + svg_lines.append('') + + return '\n'.join(svg_lines) + + def _path_to_svg_d(self, path): + """ + Convert path points to SVG d attribute + """ + if len(path) == 0: + return "" + + # Start with Move command + d_parts = [f"M {path[0][0]:.2f},{path[0][1]:.2f}"] + + # Add Line commands for remaining points + for point in path[1:]: + d_parts.append(f"L {point[0]:.2f},{point[1]:.2f}") + + return " ".join(d_parts) + + def _create_centerline_preview(self, centerline, paths, width, height): + """ + Create preview image of centerlines + """ + preview = np.zeros((height, width, 3), dtype=np.uint8) + preview.fill(255) # White background + + # Draw original centerline in light gray + preview[centerline > 127] = [200, 200, 200] + + # Draw paths in different colors + colors = [ + (255, 0, 0), # Red + (0, 255, 0), # Green + (0, 0, 255), # Blue + (255, 255, 0), # Yellow + (255, 0, 255), # Magenta + (0, 255, 255), # Cyan + ] + + for i, path in enumerate(paths): + color = colors[i % len(colors)] + path_int = path.astype(np.int32) + cv2.polylines(preview, [path_int], False, color, 2, cv2.LINE_AA) + + return preview + + def _generate_statistics(self, paths, colors, width, height, algorithm): + """ + Generate statistics about the generated paths + """ + total_length = 0 + for path in paths: + if len(path) > 1: + diffs = np.diff(path, axis=0) + lengths = np.sqrt(np.sum(diffs**2, axis=1)) + total_length += np.sum(lengths) + + avg_length = total_length / len(paths) if len(paths) > 0 else 0 + + return { + "path_count": len(paths), + "total_length": float(total_length), + "average_path_length": float(avg_length), + "image_size": [width, height], + "algorithm": algorithm, + "colors_used": len(set(colors)), + } + + +NODE_CLASS_MAPPINGS = { + "CenterlineToSVG": CenterlineToSVG +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "CenterlineToSVG": "Centerline to SVG" +} diff --git a/line_region_detector.py b/line_region_detector.py new file mode 100644 index 0000000..17b14c5 --- /dev/null +++ b/line_region_detector.py @@ -0,0 +1,296 @@ +""" +Line Region Detector Node for ComfyUI +Detects line regions and separates fill areas from line art images +""" + +import numpy as np +import torch +import cv2 +from PIL import Image +import json + + +class LineRegionDetector: + """ + Detects line regions from line art images + Separates background, lines, and fill areas + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "background_mode": (["auto", "white", "black"],), + "background_threshold": ("INT", { + "default": 240, + "min": 0, + "max": 255, + "step": 1 + }), + "line_detection_method": (["edge", "morphology", "hybrid"],), + "min_line_width": ("INT", { + "default": 1, + "min": 1, + "max": 100, + "step": 1 + }), + "max_line_width": ("INT", { + "default": 50, + "min": 1, + "max": 500, + "step": 1 + }), + "fill_handling": (["ignore", "separate", "include"],), + }, + "optional": { + "color_clustering": ("BOOLEAN", {"default": False}), + "num_colors": ("INT", { + "default": 5, + "min": 2, + "max": 20, + "step": 1 + }), + } + } + + RETURN_TYPES = ("MASK", "MASK", "IMAGE", "STRING") + RETURN_NAMES = ("line_mask", "fill_mask", "preview_image", "color_info") + FUNCTION = "detect_lines" + CATEGORY = "TJ_Vector" + + def detect_lines(self, image, background_mode="auto", background_threshold=240, + line_detection_method="hybrid", min_line_width=1, max_line_width=50, + fill_handling="separate", color_clustering=False, num_colors=5): + """ + Main function to detect line regions + """ + # Convert from ComfyUI format (B, H, W, C) to numpy + img_np = image[0].cpu().numpy() + + # Convert to 0-255 range + img_255 = (img_np * 255).astype(np.uint8) + + # Convert to grayscale + if img_255.shape[-1] >= 3: + gray = cv2.cvtColor(img_255, cv2.COLOR_RGB2GRAY) + else: + gray = img_255[..., 0] + + height, width = gray.shape + + # Step 1: Background separation + foreground_mask = self._separate_background(gray, background_mode, background_threshold) + + # Step 2: Line detection + line_mask = self._detect_line_regions( + gray, foreground_mask, line_detection_method, + min_line_width, max_line_width + ) + + # Step 3: Fill area separation + fill_mask = self._separate_fill_areas( + foreground_mask, line_mask, fill_handling + ) + + # Step 4: Color extraction (optional) + color_info = {} + if color_clustering: + color_info = self._extract_colors(img_255, line_mask, num_colors) + + # Create preview image + preview_image = self._create_preview(img_255, line_mask, fill_mask) + + # Convert masks to ComfyUI format + line_mask_tensor = torch.from_numpy(line_mask.astype(np.float32) / 255.0).unsqueeze(0) + fill_mask_tensor = torch.from_numpy(fill_mask.astype(np.float32) / 255.0).unsqueeze(0) + preview_tensor = torch.from_numpy(preview_image.astype(np.float32) / 255.0).unsqueeze(0) + + color_info_str = json.dumps(color_info, indent=2) + + return (line_mask_tensor, fill_mask_tensor, preview_tensor, color_info_str) + + def _separate_background(self, gray, mode, threshold): + """ + Separate background from foreground + """ + if mode == "auto": + # Use Otsu's method for automatic thresholding + _, foreground = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU) + elif mode == "white": + # White background + _, foreground = cv2.threshold(gray, threshold, 255, cv2.THRESH_BINARY_INV) + elif mode == "black": + # Black background + _, foreground = cv2.threshold(gray, 255 - threshold, 255, cv2.THRESH_BINARY) + + # Clean up noise + kernel = np.ones((3, 3), np.uint8) + foreground = cv2.morphologyEx(foreground, cv2.MORPH_OPEN, kernel, iterations=1) + foreground = cv2.morphologyEx(foreground, cv2.MORPH_CLOSE, kernel, iterations=1) + + return foreground + + def _detect_line_regions(self, gray, foreground_mask, method, min_width, max_width): + """ + Detect line regions using specified method + """ + if method == "edge": + return self._edge_based_detection(gray, foreground_mask, min_width, max_width) + elif method == "morphology": + return self._morphology_based_detection(foreground_mask, min_width, max_width) + elif method == "hybrid": + return self._hybrid_detection(gray, foreground_mask, min_width, max_width) + + return foreground_mask + + def _edge_based_detection(self, gray, foreground_mask, min_width, max_width): + """ + Edge-based line detection using Canny + """ + # Canny edge detection + edges = cv2.Canny(gray, 50, 150) + + # Dilate edges to create line regions + kernel_size = max(3, min_width) + kernel = np.ones((kernel_size, kernel_size), np.uint8) + dilated_edges = cv2.dilate(edges, kernel, iterations=1) + + # Combine with foreground mask + line_mask = cv2.bitwise_and(dilated_edges, foreground_mask) + + return line_mask + + def _morphology_based_detection(self, foreground_mask, min_width, max_width): + """ + Morphology-based line detection + """ + # Erode to remove fill areas + kernel_size = max(3, max_width // 2) + kernel = np.ones((kernel_size, kernel_size), np.uint8) + eroded = cv2.erode(foreground_mask, kernel, iterations=1) + + # Subtract eroded from original to get line regions + line_mask = cv2.subtract(foreground_mask, eroded) + + # Opening to clean up + kernel_small = np.ones((3, 3), np.uint8) + line_mask = cv2.morphologyEx(line_mask, cv2.MORPH_OPEN, kernel_small, iterations=1) + + return line_mask + + def _hybrid_detection(self, gray, foreground_mask, min_width, max_width): + """ + Hybrid method combining edge and morphology + """ + # Get both methods + edge_mask = self._edge_based_detection(gray, foreground_mask, min_width, max_width) + morph_mask = self._morphology_based_detection(foreground_mask, min_width, max_width) + + # Combine using OR operation + combined = cv2.bitwise_or(edge_mask, morph_mask) + + # Clean up + kernel = np.ones((3, 3), np.uint8) + combined = cv2.morphologyEx(combined, cv2.MORPH_CLOSE, kernel, iterations=1) + + return combined + + def _separate_fill_areas(self, foreground_mask, line_mask, fill_handling): + """ + Separate fill areas from lines + """ + if fill_handling == "ignore": + return np.zeros_like(foreground_mask) + + # Subtract line mask from foreground to get fill areas + fill_mask = cv2.subtract(foreground_mask, line_mask) + + if fill_handling == "separate": + # Remove small noise + kernel = np.ones((5, 5), np.uint8) + fill_mask = cv2.morphologyEx(fill_mask, cv2.MORPH_OPEN, kernel, iterations=1) + + # Remove very small regions + contours, _ = cv2.findContours(fill_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + min_area = 100 # pixels + for contour in contours: + if cv2.contourArea(contour) < min_area: + cv2.drawContours(fill_mask, [contour], -1, 0, -1) + + return fill_mask + + def _extract_colors(self, img_rgb, line_mask, num_colors): + """ + Extract color information using K-means clustering + """ + # Get pixels in line regions + mask_bool = line_mask > 127 + if not np.any(mask_bool): + return {"colors": [], "count": 0} + + line_pixels = img_rgb[mask_bool] + + if len(line_pixels) == 0: + return {"colors": [], "count": 0} + + # K-means clustering + line_pixels_float = line_pixels.astype(np.float32) + criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 100, 0.2) + k = min(num_colors, len(line_pixels)) + + if k < 1: + return {"colors": [], "count": 0} + + _, labels, centers = cv2.kmeans( + line_pixels_float, k, None, criteria, 10, cv2.KMEANS_PP_CENTERS + ) + + # Convert centers to color info + colors = [] + for i, center in enumerate(centers): + color_rgb = center.astype(int).tolist() + count = np.sum(labels == i) + colors.append({ + "rgb": color_rgb, + "hex": "#{:02x}{:02x}{:02x}".format(color_rgb[0], color_rgb[1], color_rgb[2]), + "count": int(count), + "percentage": float(count / len(labels) * 100) + }) + + # Sort by count + colors.sort(key=lambda x: x["count"], reverse=True) + + return { + "colors": colors, + "count": len(colors) + } + + def _create_preview(self, img_rgb, line_mask, fill_mask): + """ + Create preview image showing detected regions + """ + preview = img_rgb.copy() + + # Overlay line mask in red + line_overlay = np.zeros_like(preview) + line_overlay[line_mask > 127] = [255, 0, 0] # Red for lines + + # Overlay fill mask in blue + fill_overlay = np.zeros_like(preview) + fill_overlay[fill_mask > 127] = [0, 0, 255] # Blue for fills + + # Blend + preview = cv2.addWeighted(preview, 0.7, line_overlay, 0.3, 0) + preview = cv2.addWeighted(preview, 1.0, fill_overlay, 0.2, 0) + + return preview + + +NODE_CLASS_MAPPINGS = { + "LineRegionDetector": LineRegionDetector +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "LineRegionDetector": "Line Region Detector" +}