diff --git a/contourgapcloser.py b/contourgapcloser.py new file mode 100644 index 0000000..f99198f --- /dev/null +++ b/contourgapcloser.py @@ -0,0 +1,420 @@ +""" +ComfyUI Custom Node: Contour Gap Closer + +Detects broken contour lines in binary/sketch images and connects +nearby endpoints to produce closed, continuous contours. + +Algorithm: + 1. Binarize the input image (contour lines = foreground). + 2. Skeletonize to single-pixel-wide lines. + 3. Detect endpoints (pixels with exactly 1 neighbour in the skeleton). + 4. For every endpoint, find the nearest other endpoint within a + configurable max-gap distance and draw a connecting line. + 5. Optionally apply a second skeletonize pass so the output stays + single-pixel-wide. + 6. Optionally apply morphological closing first to heal very small + micro-gaps before the endpoint search. +""" + +import numpy as np +import torch +from scipy import ndimage +from scipy.spatial import KDTree +from skimage.morphology import skeletonize, disk, closing +from skimage.draw import line as sk_line + + +# ────────────────────────────────────────────── +# Helpers +# ────────────────────────────────────────────── + +def _binarize(gray: np.ndarray, threshold: float, invert: bool) -> np.ndarray: + """Return a bool array where True = contour foreground.""" + if invert: + return gray < threshold + return gray > threshold + + +def _find_endpoints(skel: np.ndarray) -> np.ndarray: + """ + Find endpoints in a skeletonised binary image. + An endpoint is a foreground pixel with exactly 1 foreground neighbour + in its 8-connected neighbourhood. + Returns an (N, 2) int array of (row, col) coordinates. + """ + # Count neighbours using convolution + kernel = np.array([[1, 1, 1], + [1, 0, 1], + [1, 1, 1]], dtype=np.uint8) + neighbour_count = ndimage.convolve(skel.astype(np.uint8), kernel, + mode='constant', cval=0) + endpoints = np.argwhere(skel & (neighbour_count == 1)) + return endpoints + + +def _direction_at_endpoint(skel: np.ndarray, pt: np.ndarray, + look_ahead: int = 6) -> np.ndarray: + """ + Estimate the tangent direction of the contour at an endpoint by + walking `look_ahead` pixels along the skeleton from `pt`. + Returns a unit vector (dr, dc). + """ + r, c = int(pt[0]), int(pt[1]) + h, w = skel.shape + visited = set() + visited.add((r, c)) + path = [(r, c)] + + for _ in range(look_ahead): + found_next = False + cr, cc = path[-1] + for dr in (-1, 0, 1): + for dc in (-1, 0, 1): + if dr == 0 and dc == 0: + continue + nr, nc = cr + dr, cc + dc + if 0 <= nr < h and 0 <= nc < w: + if skel[nr, nc] and (nr, nc) not in visited: + visited.add((nr, nc)) + path.append((nr, nc)) + found_next = True + break + if found_next: + break + if not found_next: + break + + if len(path) < 2: + return np.array([0.0, 0.0]) + + # Direction vector from last path point *towards* the endpoint + far = path[-1] + direction = np.array([pt[0] - far[0], pt[1] - far[1]], dtype=np.float64) + norm = np.linalg.norm(direction) + if norm > 0: + direction /= norm + return direction + + +def _connect_endpoints(skel: np.ndarray, max_gap: int, + direction_weight: float = 0.5, + look_ahead: int = 6) -> np.ndarray: + """ + Connect pairs of nearby endpoints in `skel`. + + For each endpoint we: + • find the nearest unmatched endpoint within `max_gap` pixels + • optionally bias toward endpoints whose tangent directions + are roughly collinear (controlled by `direction_weight`) + • draw a straight line between the matched pair + + Returns a copy of `skel` with the new connections drawn. + """ + result = skel.copy() + endpoints = _find_endpoints(skel) + + if len(endpoints) < 2: + return result + + # Pre-compute tangent directions + directions = np.array([_direction_at_endpoint(skel, ep, look_ahead) + for ep in endpoints]) + + # Build KD-tree for fast spatial lookup + tree = KDTree(endpoints) + used = set() + pairs = [] + + # Sort candidate pairs by distance so we greedily match closest first + all_dists = tree.sparse_distance_matrix(tree, max_distance=max_gap, + output_type='coo_matrix') + # Extract (i, j, dist) triples where i < j + candidates = [] + for i, j, d in zip(all_dists.row, all_dists.col, all_dists.data): + if i < j: + candidates.append((d, i, j)) + candidates.sort() + + for dist, i, j in candidates: + if i in used or j in used: + continue + + # Direction compatibility: the two tangent vectors should point + # roughly toward each other. Compute a score in [0, 1]. + d_i = directions[i] + d_j = directions[j] + if np.linalg.norm(d_i) > 0 and np.linalg.norm(d_j) > 0: + # Vector from i -> j + link = endpoints[j] - endpoints[i] + link_norm = np.linalg.norm(link) + if link_norm > 0: + link_unit = link / link_norm + # endpoint i's tangent should point toward j (dot > 0) + # endpoint j's tangent should point toward i (dot < 0 with link) + score_i = np.dot(d_i, link_unit) + score_j = np.dot(d_j, -link_unit) + compat = (score_i + score_j) / 2.0 # range [-1, 1] + else: + compat = 1.0 + else: + compat = 0.0 + + # Accept the pair if direction compatibility is above threshold + # (lower direction_weight = accept more aggressively) + threshold = -1.0 + direction_weight * 1.0 # maps [0,1] -> [-1, 0] + if compat < threshold: + continue + + used.add(i) + used.add(j) + pairs.append((i, j)) + + # Draw connecting lines + for i, j in pairs: + r0, c0 = int(endpoints[i][0]), int(endpoints[i][1]) + r1, c1 = int(endpoints[j][0]), int(endpoints[j][1]) + rr, cc = sk_line(r0, c0, r1, c1) + result[rr, cc] = True + + return result + + +# ────────────────────────────────────────────── +# ComfyUI Node +# ────────────────────────────────────────────── + +class ContourGapCloser: + """ + Joins broken contour lines to produce closed, continuous paths. + + Inputs + ------ + image : IMAGE + The sketch / contour image (white background, dark lines or vice-versa). + + Parameters + ---------- + max_gap : INT + Maximum pixel distance between two endpoints to consider joining. + threshold : FLOAT + Binarisation threshold (0-1). Pixels darker/lighter than this + are treated as contour foreground depending on `invert`. + invert : BOOL + If True, dark pixels are foreground (typical for black-on-white + sketches). If False, bright pixels are foreground. + morph_close_radius : INT + Radius of a disk structuring element for an initial morphological + closing pass. Set to 0 to skip. + re_skeletonize : BOOL + If True, skeletonize the result after gap-filling so lines stay + single-pixel-wide. + direction_weight : FLOAT + How strongly to prefer connecting endpoints whose tangent + directions are collinear (0 = ignore direction, 1 = strict). + iterations : INT + Number of detect-and-connect iterations. More iterations can + close chains of small gaps. + output_mode : ["skeleton", "dilated", "original_thickness"] + How to render the output contour lines. + dilation_radius : INT + Line thickness (radius) when output_mode is "dilated". + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "max_gap": ("INT", { + "default": 20, + "min": 1, + "max": 200, + "step": 1, + "display": "slider", + "tooltip": "Maximum pixel distance between endpoints to bridge", + }), + "threshold": ("FLOAT", { + "default": 0.5, + "min": 0.01, + "max": 0.99, + "step": 0.01, + "display": "slider", + "tooltip": "Binarization threshold (0-1)", + }), + "invert": ("BOOLEAN", { + "default": True, + "tooltip": "True = dark lines on light background", + }), + "morph_close_radius": ("INT", { + "default": 2, + "min": 0, + "max": 20, + "step": 1, + "display": "slider", + "tooltip": "Morphological closing radius to heal micro-gaps first (0=skip)", + }), + "re_skeletonize": ("BOOLEAN", { + "default": False, + "tooltip": "Skeletonize the output to single-pixel lines", + }), + "direction_weight": ("FLOAT", { + "default": 0.3, + "min": 0.0, + "max": 1.0, + "step": 0.05, + "display": "slider", + "tooltip": "Direction matching strictness (0=any, 1=strict collinear)", + }), + "iterations": ("INT", { + "default": 3, + "min": 1, + "max": 10, + "step": 1, + "tooltip": "Repeat endpoint-detection + connection N times", + }), + "output_mode": (["original_thickness", "skeleton", "dilated"],), + "dilation_radius": ("INT", { + "default": 1, + "min": 1, + "max": 10, + "step": 1, + "display": "slider", + "tooltip": "Line thickness when output_mode='dilated'", + }), + }, + } + + RETURN_TYPES = ("IMAGE", "MASK") + RETURN_NAMES = ("image", "contour_mask") + FUNCTION = "close_gaps" + CATEGORY = "image/contour" + DESCRIPTION = ("Detects broken contour lines and joins nearby endpoints " + "to produce closed, continuous paths.") + + def close_gaps(self, image: torch.Tensor, max_gap: int, + threshold: float, invert: bool, + morph_close_radius: int, re_skeletonize: bool, + direction_weight: float, iterations: int, + output_mode: str, dilation_radius: int): + + batch_out_images = [] + batch_out_masks = [] + + for b in range(image.shape[0]): + # ── 1. To numpy grayscale ── + img_np = image[b].cpu().numpy() # (H, W, C) float32 0-1 + if img_np.ndim == 3: + gray = np.mean(img_np, axis=2) + else: + gray = img_np.copy() + + # ── 2. Binarize ── + binary = _binarize(gray, threshold, invert) + + # ── 3. Optional morphological closing for micro-gaps ── + if morph_close_radius > 0: + selem = disk(morph_close_radius) + binary = closing(binary, selem) + + # ── 4. Skeletonize ── + skel = skeletonize(binary) + + # ── 5. Iterative endpoint connection ── + for _ in range(iterations): + skel = _connect_endpoints(skel, max_gap, + direction_weight=direction_weight) + + # ── 6. Optional re-skeletonize ── + if re_skeletonize: + skel = skeletonize(skel) + + # ── 7. Build output ── + if output_mode == "skeleton": + out_binary = skel + elif output_mode == "dilated": + selem = disk(dilation_radius) + out_binary = ndimage.binary_dilation(skel, structure=selem) + else: # original_thickness + # Combine original binary with the new skeleton connections + out_binary = binary | skel + + # Convert to float image (dark lines on white bg if inverted) + out_float = np.where(out_binary, 0.0, 1.0) if invert else \ + np.where(out_binary, 1.0, 0.0) + out_rgb = np.stack([out_float] * 3, axis=-1).astype(np.float32) + + mask_float = out_binary.astype(np.float32) + + batch_out_images.append(torch.from_numpy(out_rgb)) + batch_out_masks.append(torch.from_numpy(mask_float)) + + out_image = torch.stack(batch_out_images, dim=0) + out_mask = torch.stack(batch_out_masks, dim=0) + return (out_image, out_mask) + + +# ────────────────────────────────────────────── +# Debug / Visualisation Node +# ────────────────────────────────────────────── + +class ContourEndpointVisualizer: + """ + Diagnostic node: highlights detected endpoints on the image so you + can tune parameters before running the gap closer. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "threshold": ("FLOAT", { + "default": 0.5, "min": 0.01, "max": 0.99, "step": 0.01, + }), + "invert": ("BOOLEAN", {"default": True}), + "morph_close_radius": ("INT", { + "default": 2, "min": 0, "max": 20, "step": 1, + }), + "endpoint_radius": ("INT", { + "default": 5, "min": 1, "max": 20, "step": 1, + "tooltip": "Radius of the circles drawn at each endpoint", + }), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("debug_image",) + FUNCTION = "visualize" + CATEGORY = "image/contour" + DESCRIPTION = "Shows detected contour endpoints as red circles for debugging." + + def visualize(self, image, threshold, invert, morph_close_radius, + endpoint_radius): + from skimage.draw import disk as sk_disk + + batch_out = [] + for b in range(image.shape[0]): + img_np = image[b].cpu().numpy().copy() + if img_np.ndim == 2: + img_np = np.stack([img_np] * 3, axis=-1) + + gray = np.mean(img_np, axis=2) + binary = _binarize(gray, threshold, invert) + + if morph_close_radius > 0: + selem = disk(morph_close_radius) + binary = closing(binary, selem) + + skel = skeletonize(binary) + endpoints = _find_endpoints(skel) + + # Draw red circles at endpoints + h, w = img_np.shape[:2] + for ep in endpoints: + rr, cc = sk_disk((int(ep[0]), int(ep[1])), + endpoint_radius, shape=(h, w)) + img_np[rr, cc] = [1.0, 0.0, 0.0] # red + + batch_out.append(torch.from_numpy(img_np.astype(np.float32))) + + return (torch.stack(batch_out, dim=0),) \ No newline at end of file