From 6d666b6f80fdc08881b80e28107ca59ff24dbeb5 Mon Sep 17 00:00:00 2001 From: wonwoo Date: Wed, 8 Jul 2026 15:17:34 +0900 Subject: [PATCH] Add ComfyUI custom node: LSD-detected open curve smoothed via DCT Detects line segments with LSD (falling back to FLD when unavailable), chains them into a single polyline, and truncates its DCT coefficients to produce a smoothed/compressed curve, output as both an image and JSON. Co-Authored-By: Claude Sonnet 5 --- README.md | 64 +++++++++++++ __init__.py | 3 + lsd_dct_curve.py | 231 +++++++++++++++++++++++++++++++++++++++++++++++ requirements.txt | 2 + 4 files changed, 300 insertions(+) create mode 100644 README.md create mode 100644 __init__.py create mode 100644 lsd_dct_curve.py create mode 100644 requirements.txt diff --git a/README.md b/README.md new file mode 100644 index 0000000..182790e --- /dev/null +++ b/README.md @@ -0,0 +1,64 @@ +# ComfyUI-LSD-DCT-Curve + +이미지 속 개곡선(선/획/경로)을 LSD(Line Segment Detector)로 검출하고, 검출된 직선 조각들을 하나의 곡선으로 체이닝한 뒤 DCT로 스무딩/압축하는 ComfyUI 커스텀 노드입니다. + +## 설치 + +1. 이 폴더 전체를 `ComfyUI/custom_nodes/` 아래에 복사합니다. + + ``` + ComfyUI/custom_nodes/ComfyUI-LSD-DCT-Curve/ + ``` + +2. ComfyUI가 사용하는 Python 환경에서 의존성을 설치합니다. + + ``` + pip install -r ComfyUI/custom_nodes/ComfyUI-LSD-DCT-Curve/requirements.txt + ``` + + `opencv-contrib-python`은 `cv2.ximgproc.createFastLineDetector` (FLD) 폴백에 필요합니다. 이미 `opencv-python`이 설치되어 있다면 충돌을 피하기 위해 먼저 제거 후 설치하세요. + + ``` + pip uninstall opencv-python + pip install opencv-contrib-python + ``` + +3. ComfyUI를 재시작하면 노드 목록의 `curve/detection` 카테고리에 **"LSD -> DCT Curve Smoother"** 노드가 나타납니다. + +## 노드: LSD -> DCT Curve Smoother + +**입력** +- `image` (IMAGE): 개곡선이 포함된 이미지 + +**위젯** +| 이름 | 설명 | 기본값 | +|---|---|---| +| `num_dct_coeffs` | 저주파부터 남길 DCT 계수 개수. 작을수록 더 많이 스무딩/압축됨 | 16 | +| `num_samples` | DCT 전에 곡선을 균일 리샘플링할 점 개수 | 256 | +| `min_segment_length` | 이보다 짧은 LSD 조각(잡음)은 무시 | 5.0 | +| `max_gap` | 체이닝 시 허용하는 최대 끝점 간격. `0`이면 무제한(이미지 전체를 하나의 곡선으로 간주) | 0.0 | +| `canvas_mode` | 결과 곡선을 그릴 배경: `overlay_on_original` / `blank_black` / `blank_white` | overlay_on_original | +| `line_thickness` | 그려지는 곡선 두께(px) | 2 | +| `line_color` | `"R,G,B"` 형식의 곡선 색상 | 255,0,0 | + +**출력** +- `image` (IMAGE): 스무딩된 곡선이 그려진 이미지 +- `curve_json` (STRING): 검출/체이닝/DCT 결과를 담은 JSON. 필드: + - `engine`: 실제 사용된 검출 엔진 (`"LSD"` 또는 `"FLD"`) + - `num_raw_segments`: LSD/FLD가 검출한 직선 조각 개수 + - `num_chained_points`: 체이닝 후 폴리라인 점 개수 + - `num_samples`, `num_dct_coeffs_kept` + - `smoothed_curve`: 스무딩된 곡선의 `[x, y]` 좌표 배열 + - `dct_coeffs_x`, `dct_coeffs_y`: 남겨진 DCT 계수 (다른 노드에서 재구성/비교용으로 활용 가능) + +## 동작 원리 + +1. **검출**: `cv2.createLineSegmentDetector`(LSD)를 우선 시도하고, 특허 이슈로 해당 API가 없는 OpenCV 빌드에서는 `cv2.ximgproc.createFastLineDetector`(FLD)로 자동 폴백합니다. +2. **체이닝**: LSD/FLD는 곡선을 여러 개의 짧은 직선 조각으로 반환하므로, 인접한 조각의 끝점을 그리디하게 최근접 연결해 하나의 순서 있는 폴리라인으로 합칩니다. +3. **리샘플링**: DCT는 균일 샘플링을 전제하므로, 체이닝된 폴리라인을 호 길이(arc length) 기준으로 `num_samples`개 점으로 다시 뽑습니다. +4. **DCT 스무딩/압축**: x(t), y(t) 시퀀스에 각각 DCT-II를 적용하고 저주파 `num_dct_coeffs`개만 남긴 뒤 역변환해 잡음이 제거된 곡선을 복원합니다. + +## 알려진 제한 + +- 여러 개의 분리된 곡선이 한 이미지에 있는 경우, 현재 버전은 모든 조각을 하나의 곡선으로 체이닝합니다(설계상 가정). 분리된 곡선을 각각 다루려면 `max_gap`을 낮추거나 별도 배치 처리가 필요합니다. +- DCT 계수를 적게 남길수록 곡선 양 끝단에 약간의 처짐/리플이 생길 수 있습니다(경계 처리 특성). diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..8f9c02d --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .lsd_dct_curve import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/lsd_dct_curve.py b/lsd_dct_curve.py new file mode 100644 index 0000000..66da00b --- /dev/null +++ b/lsd_dct_curve.py @@ -0,0 +1,231 @@ +import json +from collections import deque + +import cv2 +import numpy as np +import torch + + +def _tensor_to_gray(img_tensor): + """img_tensor: (H, W, C) float32 in [0,1] -> (uint8 RGB image, uint8 gray image).""" + img = np.ascontiguousarray((img_tensor.cpu().numpy() * 255.0).clip(0, 255).astype(np.uint8)) + if img.shape[-1] == 1: + gray = img[..., 0] + else: + gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY) + return img, gray + + +def _detect_segments(gray): + """Return (Nx4 float32 array of x1,y1,x2,y2, engine name). Tries classic LSD first, + falls back to FLD since cv2.createLineSegmentDetector is patent-restricted and + missing/broken on some OpenCV builds.""" + lines = None + engine = None + try: + lsd = cv2.createLineSegmentDetector(cv2.LSD_REFINE_STD) + lines = lsd.detect(gray)[0] + engine = "LSD" + except (AttributeError, cv2.error): + lines = None + + if lines is None: + if not hasattr(cv2, "ximgproc"): + raise RuntimeError( + "cv2.createLineSegmentDetector is unavailable on this OpenCV build and " + "cv2.ximgproc is missing for the FLD fallback. Install opencv-contrib-python." + ) + fld = cv2.ximgproc.createFastLineDetector() + lines = fld.detect(gray) + engine = "FLD" + + if lines is None: + return np.zeros((0, 4), dtype=np.float32), engine or "none" + return lines.reshape(-1, 4).astype(np.float32), engine + + +def _chain_segments(segs, min_length=5.0, max_gap=0.0): + """Greedily chain disjoint line segments end-to-end into a single ordered polyline, + always extending from whichever open end (head or tail) is closest to an + unused segment's endpoint. Assumes the image contains one dominant curve.""" + segs = [s for s in segs if np.hypot(s[2] - s[0], s[3] - s[1]) >= min_length] + if not segs: + return np.zeros((0, 2), dtype=np.float32) + + segs = sorted(segs, key=lambda s: np.hypot(s[2] - s[0], s[3] - s[1]), reverse=True) + starts = [np.array([s[0], s[1]], dtype=np.float32) for s in segs] + ends = [np.array([s[2], s[3]], dtype=np.float32) for s in segs] + + n = len(segs) + used = [False] * n + used[0] = True + chain = deque([starts[0], ends[0]]) + remaining = n - 1 + + while remaining > 0: + head, tail = chain[0], chain[-1] + best = None # (dist, idx, end, flip) + for i in range(n): + if used[i]: + continue + p1, p2 = starts[i], ends[i] + for dist, end, flip in ( + (np.linalg.norm(p1 - tail), "tail", False), + (np.linalg.norm(p2 - tail), "tail", True), + (np.linalg.norm(p1 - head), "head", True), + (np.linalg.norm(p2 - head), "head", False), + ): + if best is None or dist < best[0]: + best = (dist, i, end, flip) + + dist, idx, end, flip = best + if max_gap > 0.0 and dist > max_gap: + break + + used[idx] = True + p1, p2 = starts[idx], ends[idx] + if end == "tail": + if flip: + chain.append(p2); chain.append(p1) + else: + chain.append(p1); chain.append(p2) + else: + if flip: + chain.appendleft(p1); chain.appendleft(p2) + else: + chain.appendleft(p2); chain.appendleft(p1) + remaining -= 1 + + return np.array(chain, dtype=np.float32) + + +def _resample_polyline(points, num_samples): + """Resample a polyline to num_samples evenly spaced points by arc length, + which DCT smoothing requires (it assumes uniform sampling).""" + if len(points) < 2: + return np.zeros((0, 2), dtype=np.float32) + deltas = np.diff(points, axis=0) + seg_lengths = np.hypot(deltas[:, 0], deltas[:, 1]) + cumlen = np.concatenate([[0.0], np.cumsum(seg_lengths)]) + total_len = cumlen[-1] + if total_len <= 1e-6: + return np.repeat(points[:1], num_samples, axis=0) + sample_dist = np.linspace(0, total_len, num_samples) + xs = np.interp(sample_dist, cumlen, points[:, 0]) + ys = np.interp(sample_dist, cumlen, points[:, 1]) + return np.stack([xs, ys], axis=1).astype(np.float32) + + +def _dct_smooth(coords_1d, num_coeffs): + """Keep the first num_coeffs DCT-II coefficients (low frequencies) and + reconstruct via inverse DCT -> smoothed / compressed curve axis.""" + src = coords_1d.reshape(-1, 1).astype(np.float32) + coeffs = cv2.dct(src) + kept = coeffs.copy() + kept[num_coeffs:] = 0 + smoothed = cv2.idct(kept).flatten() + return smoothed, coeffs.flatten() + + +def _parse_color(color_str): + try: + r, g, b = (int(v.strip()) for v in color_str.split(",")) + return (max(0, min(255, r)), max(0, min(255, g)), max(0, min(255, b))) + except Exception: + return (255, 0, 0) + + +def _draw_curve(base_rgb_uint8, points, mode, color, thickness): + if mode == "blank_black": + canvas = np.zeros_like(base_rgb_uint8) + elif mode == "blank_white": + canvas = np.full_like(base_rgb_uint8, 255) + else: + canvas = base_rgb_uint8.copy() + if len(points) >= 2: + pts = points.round().astype(np.int32).reshape(-1, 1, 2) + cv2.polylines(canvas, [pts], isClosed=False, color=color, thickness=thickness, lineType=cv2.LINE_AA) + return canvas + + +class LSDDCTCurveNode: + """Detects open-curve line structure in an image with LSD (falling back to FLD), + chains the resulting line segments into a single ordered polyline, and smooths / + compresses it by truncating its DCT coefficients.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "num_dct_coeffs": ("INT", {"default": 16, "min": 2, "max": 512, "step": 1}), + "num_samples": ("INT", {"default": 256, "min": 16, "max": 4096, "step": 1}), + "min_segment_length": ("FLOAT", {"default": 5.0, "min": 0.0, "max": 1000.0, "step": 0.5}), + "max_gap": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 5000.0, "step": 1.0}), + "canvas_mode": (["overlay_on_original", "blank_black", "blank_white"],), + "line_thickness": ("INT", {"default": 2, "min": 1, "max": 20, "step": 1}), + "line_color": ("STRING", {"default": "255,0,0"}), + } + } + + RETURN_TYPES = ("IMAGE", "STRING") + RETURN_NAMES = ("image", "curve_json") + FUNCTION = "process" + CATEGORY = "curve/detection" + + def process( + self, + image, + num_dct_coeffs, + num_samples, + min_segment_length, + max_gap, + canvas_mode, + line_thickness, + line_color, + ): + num_dct_coeffs = min(num_dct_coeffs, num_samples) + color = _parse_color(line_color) + + out_images = [] + out_records = [] + + for b in range(image.shape[0]): + rgb_uint8, gray = _tensor_to_gray(image[b]) + segs, engine = _detect_segments(gray) + chain_pts = _chain_segments(segs, min_length=min_segment_length, max_gap=max_gap) + + if len(chain_pts) >= 2: + resampled = _resample_polyline(chain_pts, num_samples) + smoothed_x, dct_x = _dct_smooth(resampled[:, 0], num_dct_coeffs) + smoothed_y, dct_y = _dct_smooth(resampled[:, 1], num_dct_coeffs) + smoothed_pts = np.stack([smoothed_x, smoothed_y], axis=1) + else: + smoothed_pts = np.zeros((0, 2), dtype=np.float32) + dct_x = dct_y = np.zeros((0,), dtype=np.float32) + + canvas = _draw_curve(rgb_uint8, smoothed_pts, canvas_mode, color, line_thickness) + out_images.append(torch.from_numpy(canvas.astype(np.float32) / 255.0)) + + out_records.append({ + "engine": engine, + "num_raw_segments": int(len(segs)), + "num_chained_points": int(len(chain_pts)), + "num_samples": int(num_samples), + "num_dct_coeffs_kept": int(num_dct_coeffs), + "smoothed_curve": np.round(smoothed_pts, 3).tolist(), + "dct_coeffs_x": np.round(dct_x[:num_dct_coeffs], 4).tolist(), + "dct_coeffs_y": np.round(dct_y[:num_dct_coeffs], 4).tolist(), + }) + + out_batch = torch.stack(out_images, dim=0) + return (out_batch, json.dumps(out_records)) + + +NODE_CLASS_MAPPINGS = { + "LSDDCTCurve": LSDDCTCurveNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "LSDDCTCurve": "LSD -> DCT Curve Smoother", +} diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..4402a9f --- /dev/null +++ b/requirements.txt @@ -0,0 +1,2 @@ +opencv-contrib-python +numpy