Update nodes with new version
This commit is contained in:
unknown
@@ -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",
|
||||
}
|
||||
+395
@@ -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))
|
||||
@@ -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),)
|
||||
Reference in New Issue
Block a user