From a112b4dc159fc4a2f4ce6ce469be39c3aeae07fc Mon Sep 17 00:00:00 2001 From: unknown Date: Sun, 22 Mar 2026 20:18:10 +0700 Subject: [PATCH] Update nodes with new version --- __init__.py | 3 + laserpath.py | 395 +++++++++++++++++++++++++++++++++++++++++++++++++++ ocr.py | 121 ---------------- 3 files changed, 398 insertions(+), 121 deletions(-) create mode 100644 laserpath.py delete mode 100644 ocr.py diff --git a/__init__.py b/__init__.py index 2c1685d..f3fbdc6 100644 --- a/__init__.py +++ b/__init__.py @@ -11,6 +11,7 @@ from .autocrop import * from .numfunction import * from .linecounter import * from .numstrcleanup import * +from .laserpath import * NODE_CLASS_MAPPINGS = { "RasterToUniformSVG": RasterToSVGConverter, @@ -29,6 +30,7 @@ NODE_CLASS_MAPPINGS = { "RoundDownNode": RoundDownNode, "CSVValueExtractor": CSVValueExtractor, "NumberStringCleanup": NumberStringCleanup, + "LaserPathTracer": LaserPathTracerNode, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -48,4 +50,5 @@ NODE_DISPLAY_NAME_MAPPINGS = { "RoundDownNode": "Round Down", "CSVValueExtractor": "CSV Value Extractor", "NumberStringCleanup": "Number String Cleanup", + "LaserPathTracer": "Laser Cutter Path Tracer", } \ No newline at end of file diff --git a/laserpath.py b/laserpath.py new file mode 100644 index 0000000..e5626d6 --- /dev/null +++ b/laserpath.py @@ -0,0 +1,395 @@ +""" +ComfyUI Custom Node: Laser Cutter Path Tracer +Detects internal contours (holes) in a B&W diagram and traces the optimal +laser cutter head path between them using nearest-neighbor + 2-opt TSP. +""" + +import numpy as np +import cv2 +try: + import torch +except ImportError: + torch = None +from itertools import combinations + + +# ────────────────────────────────────────────────────────────────────── +# Geometry helpers +# ────────────────────────────────────────────────────────────────────── + +def contour_centroid(contour): + """Return (cx, cy) of a contour using image moments.""" + M = cv2.moments(contour) + if M["m00"] == 0: + # fallback: mean of all points + pts = contour.reshape(-1, 2) + return float(pts[:, 0].mean()), float(pts[:, 1].mean()) + return float(M["m10"] / M["m00"]), float(M["m01"] / M["m00"]) + + +def contour_nearest_edge_point(contour, ref_point): + """Return the point on `contour` closest to `ref_point`.""" + pts = contour.reshape(-1, 2).astype(np.float64) + ref = np.array(ref_point, dtype=np.float64) + dists = np.linalg.norm(pts - ref, axis=1) + return tuple(pts[np.argmin(dists)].astype(int)) + + +def euclidean(a, b): + return ((a[0] - b[0]) ** 2 + (a[1] - b[1]) ** 2) ** 0.5 + + +# ────────────────────────────────────────────────────────────────────── +# TSP solvers +# ────────────────────────────────────────────────────────────────────── + +def nearest_neighbor_open(points, start_idx, end_idx): + """ + Greedy nearest-neighbor for an open path. + Starts at `start_idx`, must finish at `end_idx`. + Returns ordered list of indices. + """ + n = len(points) + if n <= 2: + return list(range(n)) + + visited = [False] * n + path = [start_idx] + visited[start_idx] = True + # Reserve end point + visited[end_idx] = True + + current = start_idx + remaining = n - 2 # exclude start and end + + for _ in range(remaining): + best_dist = float("inf") + best_idx = -1 + for j in range(n): + if not visited[j]: + d = euclidean(points[current], points[j]) + if d < best_dist: + best_dist = d + best_idx = j + if best_idx == -1: + break + visited[best_idx] = True + path.append(best_idx) + current = best_idx + + path.append(end_idx) + return path + + +def two_opt_improve(path, points, max_iterations=1000): + """ + 2-opt local search improvement for an open path. + Keeps first and last elements fixed. + """ + def path_length(p): + return sum(euclidean(points[p[i]], points[p[i + 1]]) for i in range(len(p) - 1)) + + best_dist = path_length(path) + improved = True + iteration = 0 + + while improved and iteration < max_iterations: + improved = False + iteration += 1 + # Only reverse interior segments (keep index 0 and -1 fixed) + for i in range(1, len(path) - 2): + for j in range(i + 1, len(path) - 1): + new_path = path[:i] + path[i : j + 1][::-1] + path[j + 1 :] + new_dist = path_length(new_path) + if new_dist < best_dist - 1e-6: + path = new_path + best_dist = new_dist + improved = True + break + if improved: + break + + return path + + +def solve_tsp(points, start_idx, end_idx): + """Solve open-path TSP with fixed start and end using NN + 2-opt.""" + if len(points) <= 1: + return list(range(len(points))) + if len(points) == 2: + return [start_idx, end_idx] + + path = nearest_neighbor_open(points, start_idx, end_idx) + path = two_opt_improve(path, points) + return path + + +# ────────────────────────────────────────────────────────────────────── +# Contour classification +# ────────────────────────────────────────────────────────────────────── + +def find_internal_contours(binary_img, min_area=30, max_area_ratio=0.5): + """ + Detect internal contours (holes) inside the external boundary. + + Strategy: + 1. Find all contours with full hierarchy. + 2. Identify the external contour as the largest by area. + 3. Internal contours are contours whose centroid lies inside + the external contour and whose area is much smaller. + + Returns: (external_contour, list_of_internal_contours) + """ + contours, hierarchy = cv2.findContours( + binary_img, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE + ) + + if not contours: + return None, [] + + # Sort by area descending + areas = [cv2.contourArea(c) for c in contours] + max_area = max(areas) if areas else 0 + + if max_area == 0: + return None, [] + + # External contour = largest + ext_idx = int(np.argmax(areas)) + external_contour = contours[ext_idx] + ext_area = areas[ext_idx] + + # Internal contours: smaller contours whose centroid is inside the external + candidates = [] + for i, c in enumerate(contours): + if i == ext_idx: + continue + a = areas[i] + if a < min_area: + continue + if a > ext_area * max_area_ratio: + continue + + cx, cy = contour_centroid(c) + # Check if centroid is inside external contour + inside = cv2.pointPolygonTest(external_contour, (cx, cy), False) + if inside >= 0: + candidates.append((c, cx, cy, a)) + + # Deduplicate: drawn circle strokes produce two contours (inner/outer) + # with nearly identical centroids. Keep the larger one per cluster. + dedup_radius = 15 # px — merge contours with centroids this close + candidates.sort(key=lambda x: -x[3]) # sort by area descending + internal = [] + used_centroids = [] + for c, cx, cy, a in candidates: + duplicate = False + for ux, uy in used_centroids: + if ((cx - ux) ** 2 + (cy - uy) ** 2) ** 0.5 < dedup_radius: + duplicate = True + break + if not duplicate: + internal.append(c) + used_centroids.append((cx, cy)) + + return external_contour, internal + + +# ────────────────────────────────────────────────────────────────────── +# Drawing +# ────────────────────────────────────────────────────────────────────── + +def draw_path( + image, + ordered_centroids, + line_color=(80, 80, 80), + start_color=(0, 180, 0), + end_color=(0, 0, 220), + point_color=(60, 60, 60), + line_thickness=2, + point_radius=6, +): + """Draw the traced path on the image.""" + out = image.copy() + if len(out.shape) == 2: + out = cv2.cvtColor(out, cv2.COLOR_GRAY2BGR) + + n = len(ordered_centroids) + if n == 0: + return out + + # Draw lines + for i in range(n - 1): + pt1 = (int(ordered_centroids[i][0]), int(ordered_centroids[i][1])) + pt2 = (int(ordered_centroids[i + 1][0]), int(ordered_centroids[i + 1][1])) + cv2.line(out, pt1, pt2, line_color, line_thickness, cv2.LINE_AA) + + # Draw points + for i, (cx, cy) in enumerate(ordered_centroids): + pt = (int(cx), int(cy)) + if i == 0: + cv2.circle(out, pt, point_radius + 2, start_color, -1, cv2.LINE_AA) + elif i == n - 1: + cv2.circle(out, pt, point_radius + 2, end_color, -1, cv2.LINE_AA) + else: + cv2.circle(out, pt, point_radius, point_color, -1, cv2.LINE_AA) + + return out + + +# ────────────────────────────────────────────────────────────────────── +# ComfyUI Node +# ────────────────────────────────────────────────────────────────────── + +class LaserPathTracerNode: + """ + Traces the optimal laser cutter path between internal holes in a + B&W diagram. Start = bottom-right, End = bottom-left. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "binary_threshold": ( + "INT", + {"default": 128, "min": 0, "max": 255, "step": 1}, + ), + "min_hole_area": ( + "INT", + {"default": 30, "min": 1, "max": 10000, "step": 1}, + ), + "max_hole_area_pct": ( + "FLOAT", + {"default": 0.4, "min": 0.01, "max": 0.99, "step": 0.01}, + ), + "line_thickness": ( + "INT", + {"default": 2, "min": 1, "max": 10, "step": 1}, + ), + "point_radius": ( + "INT", + {"default": 6, "min": 2, "max": 20, "step": 1}, + ), + "invert_image": ( + "BOOLEAN", + {"default": False}, + ), + }, + } + + RETURN_TYPES = ("IMAGE", "STRING") + RETURN_NAMES = ("traced_image", "path_info") + FUNCTION = "trace_path" + CATEGORY = "Laser/Path" + + def trace_path( + self, + image, + binary_threshold, + min_hole_area, + max_hole_area_pct, + line_thickness, + point_radius, + invert_image, + ): + # ── 1. Convert ComfyUI image (B,H,W,C float 0-1) to OpenCV ── + img_np = image[0].cpu().numpy() # (H, W, C) float 0-1 + img_uint8 = (img_np * 255).clip(0, 255).astype(np.uint8) + + if img_uint8.shape[2] == 4: + img_bgr = cv2.cvtColor(img_uint8, cv2.COLOR_RGBA2BGR) + else: + img_bgr = cv2.cvtColor(img_uint8, cv2.COLOR_RGB2BGR) + + gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY) + + # ── 2. Binarize ── + _, binary = cv2.threshold(gray, binary_threshold, 255, cv2.THRESH_BINARY) + if invert_image: + binary = cv2.bitwise_not(binary) + + # For contour detection we need black shapes on white background. + # Contours are found on white regions in RETR_TREE with inverted image, + # or we detect on the inverse. Let's detect contours of the black lines: + binary_inv = cv2.bitwise_not(binary) + + # ── 3. Detect contours ── + external_contour, internal_contours = find_internal_contours( + binary_inv, + min_area=min_hole_area, + max_area_ratio=max_hole_area_pct, + ) + + h, w = gray.shape[:2] + + if not internal_contours: + # Fallback: try detecting on the non-inverted binary + # (holes as white regions inside a black border) + external_contour, internal_contours = find_internal_contours( + binary, + min_area=min_hole_area, + max_area_ratio=max_hole_area_pct, + ) + + if not internal_contours: + info = "No internal contours (holes) detected. Try adjusting threshold or min_hole_area." + out_tensor = torch.from_numpy(img_np).unsqueeze(0) + return (out_tensor, info) + + # ── 4. Compute centroids ── + centroids = [contour_centroid(c) for c in internal_contours] + + # ── 5. Add virtual start (bottom-right) and end (bottom-left) ── + # We snap these to the nearest hole edge if within range, + # otherwise use corner coordinates offset slightly inward. + margin_x = int(w * 0.02) + margin_y = int(h * 0.02) + start_ref = (w - margin_x, h - margin_y) # bottom-right + end_ref = (margin_x, h - margin_y) # bottom-left + + all_points = [start_ref] + centroids + [end_ref] + start_idx = 0 + end_idx = len(all_points) - 1 + + # ── 6. Solve TSP ── + path_order = solve_tsp(all_points, start_idx, end_idx) + ordered_points = [all_points[i] for i in path_order] + + # ── 7. Draw ── + result_bgr = draw_path( + img_bgr, + ordered_points, + line_thickness=line_thickness, + point_radius=point_radius, + ) + + # ── 8. Convert back to ComfyUI format ── + result_rgb = cv2.cvtColor(result_bgr, cv2.COLOR_BGR2RGB) + result_float = result_rgb.astype(np.float32) / 255.0 + out_tensor = torch.from_numpy(result_float).unsqueeze(0) + + # ── 9. Build info string ── + total_dist = sum( + euclidean(ordered_points[i], ordered_points[i + 1]) + for i in range(len(ordered_points) - 1) + ) + info_lines = [ + f"Holes detected: {len(internal_contours)}", + f"Path points: {len(ordered_points)} (incl. start & end)", + f"Total path length: {total_dist:.1f} px", + f"Start: ({ordered_points[0][0]:.0f}, {ordered_points[0][1]:.0f}) [bottom-right]", + f"End: ({ordered_points[-1][0]:.0f}, {ordered_points[-1][1]:.0f}) [bottom-left]", + "", + "Visit order (x, y):", + ] + for i, pt in enumerate(ordered_points): + label = "" + if i == 0: + label = " <- START" + elif i == len(ordered_points) - 1: + label = " <- END" + info_lines.append(f" {i}: ({pt[0]:.0f}, {pt[1]:.0f}){label}") + + return (out_tensor, "\n".join(info_lines)) \ No newline at end of file diff --git a/ocr.py b/ocr.py deleted file mode 100644 index 73c8c58..0000000 --- a/ocr.py +++ /dev/null @@ -1,121 +0,0 @@ -import re -import numpy as np -from PIL import Image - - -class MaxValuePaddleOCR: - """ - A ComfyUI node that uses PaddleOCR to extract numbers from an image - and returns the largest numeric value found. - """ - - def __init__(self): - self.ocr = None - - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "image": ("IMAGE",), - }, - "optional": { - "use_gpu": ("BOOLEAN", {"default": False}), - "lang": (["en", "ch", "japan", "korean", "french", "german"], {"default": "en"}), - } - } - - RETURN_TYPES = ("FLOAT",) - RETURN_NAMES = ("max_value",) - FUNCTION = "extract_max_number" - CATEGORY = "image/text" - - def _init_ocr(self, use_gpu=False, lang="en"): - """Initialize PaddleOCR if not already done.""" - if self.ocr is None: - from paddleocr import PaddleOCR - self.ocr = PaddleOCR( - use_angle_cls=True, - lang=lang, - use_gpu=use_gpu, - show_log=False - ) - return self.ocr - - def _extract_numbers(self, text): - """ - Extract all numeric values (integers and floats) from text. - Handles various number formats including negative numbers and decimals. - """ - # Pattern matches integers and floats, including negative numbers - # Examples: 123, -456, 3.14, -2.5, .5, -.5 - pattern = r'-?\d+\.?\d*|-?\.\d+' - - numbers = [] - matches = re.findall(pattern, text) - - for match in matches: - try: - # Try to convert to float (works for both int and float strings) - num = float(match) - numbers.append(num) - except ValueError: - continue - - return numbers - - def extract_max_number(self, image, use_gpu=False, lang="en"): - """ - Main function to extract the maximum number from an image. - - Args: - image: ComfyUI IMAGE tensor (B, H, W, C) in range [0, 1] - use_gpu: Whether to use GPU for OCR - lang: Language for OCR recognition - - Returns: - Tuple containing the maximum numeric value found, or 0.0 if none found - """ - # Initialize OCR - ocr = self._init_ocr(use_gpu=use_gpu, lang=lang) - - # Convert ComfyUI image tensor to numpy array - # ComfyUI images are (B, H, W, C) with values in [0, 1] - if isinstance(image, np.ndarray): - img_array = image - else: - # Convert from torch tensor if needed - img_array = image.cpu().numpy() - - # Handle batch dimension - process first image - if len(img_array.shape) == 4: - img_array = img_array[0] - - # Convert from [0, 1] to [0, 255] and ensure uint8 - img_array = (img_array * 255).astype(np.uint8) - - # Convert to PIL Image then back to numpy for PaddleOCR - # PaddleOCR expects BGR or RGB numpy array - pil_image = Image.fromarray(img_array) - img_for_ocr = np.array(pil_image) - - # Run OCR - result = ocr.ocr(img_for_ocr, cls=True) - - # Extract all numbers from OCR results - all_numbers = [] - - if result and result[0]: - for line in result[0]: - if line and len(line) >= 2: - # line[1] contains (text, confidence) - text = line[1][0] if isinstance(line[1], (list, tuple)) else str(line[1]) - numbers = self._extract_numbers(text) - all_numbers.extend(numbers) - - # Return the maximum value, or 0.0 if no numbers found - if all_numbers: - max_value = max(all_numbers) - else: - max_value = 0.0 - - return (float(max_value),) \ No newline at end of file