diff --git a/DetectInnerBox.py b/DetectInnerBox.py new file mode 100644 index 0000000..4707e72 --- /dev/null +++ b/DetectInnerBox.py @@ -0,0 +1,48 @@ +import cv2 +import torch +import numpy as np + + +class DetectInnerBox: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mask": ("MASK",), # 只接 MASK + "connectivity": ("INT", {"default": 4, "min": 4, "max": 8}), + } + } + + RETURN_TYPES = ("INT", "INT") # width, height + RETURN_NAMES = ("inner_w", "inner_h") + FUNCTION = "run" + CATEGORY = "DrawTools" + + # ---------------- 主逻辑 ---------------- # + def run(self, mask, connectivity=4): + # ① tensor → numpy 单通道 0/255 + if isinstance(mask, list): # ComfyUI 有时包在 list 里 + mask = mask[0] + if mask.ndim == 3: # [1,H,W] → [H,W] + mask = mask[0] + np_mask = mask.cpu().numpy().astype("uint8") * 255 # 0/255 + + H, W = np_mask.shape[:2] + + # ② flood-fill 把连边的白色清零 + pad = cv2.copyMakeBorder(np_mask, 1, 1, 1, 1, cv2.BORDER_CONSTANT, value=0) + flags = 4 if connectivity == 4 else 8 + cv2.floodFill( + pad, None, (0, 0), 0, loDiff=0, upDiff=0, flags=flags | (255 << 8) + ) + inner = pad[1:-1, 1:-1] + + # ③ 取剩余像素的 bbox + ys, xs = np.where(inner == 255) + if ys.size == 0: + return 0, 0 + + y0, y1 = ys.min(), ys.max() + x0, x1 = xs.min(), xs.max() + + return int(x1 - x0 + 1), int(y1 - y0 + 1) diff --git a/PasteIntoFrame.py b/PasteIntoFrame.py new file mode 100644 index 0000000..ec8f654 --- /dev/null +++ b/PasteIntoFrame.py @@ -0,0 +1,113 @@ +import cv2 +import torch +import numpy as np + + +class PasteIntoFrame: + """ + 根据 mask 中央矩形尺寸缩放 content,再用 *原始 mask* 把内容贴进 frame。 + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "frame_image": ("IMAGE",), + "content_image": ("IMAGE",), + "mask_image": ("IMAGE",), + "threshold": ( + "FLOAT", + {"default": 0.2, "min": 0.0, "max": 1.0, "step": 0.01}, + ), + "connectivity": ("INT", {"default": 4, "min": 4, "max": 8}), + "keep_aspect": ("BOOLEAN", {"default": False}), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("composited_image",) + FUNCTION = "run" + CATEGORY = "DrawTools" + + # ---------------- 主逻辑 ---------------- # + def run( + self, + frame_image, + content_image, + mask_image, + threshold=0.2, + connectivity=4, + keep_aspect=False, + ): + frame = self._tensor_to_np(frame_image) # H,W,3 或 4 + content = self._tensor_to_np(content_image) # Hc,Wc,3 + mask = self._tensor_to_np( + mask_image, + gray=True, # 0/255 + thr=threshold, + ) + + H, W = mask.shape + + # ① 生成 inner mask 仅用于 bbox 计算 + pad = cv2.copyMakeBorder(mask, 1, 1, 1, 1, cv2.BORDER_CONSTANT, value=0) + flags = 4 if connectivity == 4 else 8 + cv2.floodFill(pad, None, (0, 0), 0, flags=flags | (255 << 8)) + inner = pad[1:-1, 1:-1] + + ys, xs = np.where(inner == 255) + if ys.size == 0: + raise ValueError("mask 中央区域为空!") + + y0, y1 = ys.min(), ys.max() + x0, x1 = xs.min(), xs.max() + target_w, target_h = x1 - x0 + 1, y1 - y0 + 1 + + # ② 把 content 缩放/填充到和 frame 同尺寸的 canvas + canvas = frame[:, :, :3].copy() # 复制框的 RGB (不含 alpha) + if keep_aspect: + h0, w0 = content.shape[:2] + scale = min(target_w / w0, target_h / h0) + new_w, new_h = int(w0 * scale), int(h0 * scale) + resized = cv2.resize(content, (new_w, new_h), interpolation=cv2.INTER_AREA) + y_off = y0 + (target_h - new_h) // 2 + x_off = x0 + (target_w - new_w) // 2 + canvas[y_off : y_off + new_h, x_off : x_off + new_w] = resized + else: + resized = cv2.resize( + content, (target_w, target_h), interpolation=cv2.INTER_AREA + ) + canvas[y0 : y1 + 1, x0 : x1 + 1] = resized + + # ③ 用 **原始 mask** 贴图 + result = frame[:, :, :3].copy() + cv2.copyTo(canvas, mask, result) # mask=原始 => 保留全部异形区域 + + # ④ 恢复透明度(如果 frame 有 α) + if frame.shape[2] == 4: + alpha = frame[:, :, 3:] + result = np.concatenate([result, alpha], axis=2) + + return (self._np_to_tensor(result),) + + # --------- tensor ↔ numpy 工具 --------- # + def _tensor_to_np(self, img, gray=False, thr=0.2): + if isinstance(img, list): + img = img[0] + if img.ndim == 3 and img.shape[0] <= 4: # [C,H,W] → [H,W,C] + img = img.permute(1, 2, 0) + arr = img.cpu().numpy() + if arr.dtype != np.bool_ and arr.max() <= 1.0: + arr = arr * 255.0 + arr = arr.astype("uint8") + if gray: + if arr.ndim == 3: + arr = cv2.cvtColor(arr, cv2.COLOR_BGR2GRAY) + arr = (arr > int(thr * 255)).astype("uint8") * 255 + return arr + + def _np_to_tensor(self, arr): + if arr.ndim == 2: + arr = arr[:, :, None] + tensor = torch.from_numpy(arr.astype("float32") / 255.0).permute(2, 0, 1) + return tensor.unsqueeze(0) diff --git a/README.md b/README.md index f155fa0..1e8965d 100644 --- a/README.md +++ b/README.md @@ -1,47 +1,54 @@ -# YGO Tools for ComfyUI +# Draw Tools for ComfyUI -一组针对游戏王卡图制作的ComfyUI自定义节点集合。 +一组针对图像处理的ComfyUI自定义节点集合。 ## 功能节点 -### Tile Image +### Detect Inner Box -- 将输入图像水平平铺指定次数(最多12个) -- 可控制图像之间的间距(基于图像宽度的比例) -- 支持左对齐或右对齐排列 -- 输出带有完整mask的图像 +- 检测掩码中不与边缘相连的白色区域 +- 返回区域最小外接矩形的宽度和高度 +- 支持设置阈值和连通性参数 +- 适用于检测图像中孤立区域的尺寸 -### Draw Line +### Paste Into Frame -- 绘制自定义线条 -- 支持设置线条长度和宽度 -- 使用HEX格式设置颜色(如#FFFFFF) -- 支持设置透明度(如#FFFFFFFF) -- 支持任意角度(0-360度)的线条绘制 +- 根据掩码中央矩形尺寸缩放内容图像 +- 使用原始掩码将内容贴进框架图像 +- 支持保持内容图像的原始宽高比 +- 自动处理透明度(如果框架图像有alpha通道) ## 安装 ```bash cd custom_nodes -git clone https://github.com/Ky11le/ygo_tools +git clone https://github.com/Ky11le/draw_tools +cd draw_tools +pip install -r requirements.txt ``` +## 依赖要求 + +本插件需要以下依赖项: +- opencv-python >= 4.5.0 +- torch >= 1.7.0 +- numpy >= 1.19.0 + ## 使用方法 -节点会在ComfyUI的"ygo_tools"类别下出现。 +节点会在ComfyUI的"DrawTools"类别下出现。 -### Tile Image 参数 +### Detect Inner Box 参数 -- `image`: 输入图像 -- `mask`: 输入mask -- `num_tiles`: 要平铺的图像数量(1-12) -- `spacing_factor`: 图像间距系数(相对于图像宽度) -- `align`: 对齐方式(left/right) +- `mask_image`: 输入掩码图像 +- `threshold`: 二值化阈值(0.0-1.0) +- `connectivity`: 连通性(4或8) -### Draw Line 参数 +### Paste Into Frame 参数 -- `length`: 线条长度 -- `width`: 线条宽度 -- `color`: HEX颜色值(如#FFFFFF) -- `angle`: 线条角度(0-360度) -- `padding`: 边距大小 \ No newline at end of file +- `frame_image`: 框架图像(将作为背景) +- `content_image`: 内容图像(将被贴入框架) +- `mask_image`: 掩码图像(决定贴图区域) +- `threshold`: 二值化阈值(0.0-1.0) +- `connectivity`: 连通性(4或8) +- `keep_aspect`: 是否保持内容图像的原始宽高比 \ No newline at end of file diff --git a/__init__.py b/__init__.py index 06d5510..1e45558 100644 --- a/__init__.py +++ b/__init__.py @@ -1,12 +1,12 @@ -from .tile_image import TileImage -from .draw_line import DrawLine +from .DetectInnerBox import DetectInnerBox +from .PasteIntoFrame import PasteIntoFrame NODE_CLASS_MAPPINGS = { - "TileImage": TileImage, - "DrawLine": DrawLine + "DetectInnerBox": DetectInnerBox, + "PasteIntoFrame": PasteIntoFrame, } NODE_DISPLAY_NAME_MAPPINGS = { - "TileImage": "Tile Image", - "DrawLine": "Draw Line" -} \ No newline at end of file + "DetectInnerBox": "Detect Inner Box", + "PasteIntoFrame": "Paste Into Frame", +} diff --git a/draw_line.py b/draw_line.py deleted file mode 100644 index 82a22b6..0000000 --- a/draw_line.py +++ /dev/null @@ -1,96 +0,0 @@ -import torch -import numpy as np -import math - -class DrawLine: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "length": ("INT", {"default": 512, "min": 1, "max": 4096}), - "width": ("INT", {"default": 8, "min": 1, "max": 1024}), - "color": ("STRING", {"default": "#FFFFFF", "multiline": False}), - "angle": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 360.0, "step": 1.0}), - "padding": ("INT", {"default": 0, "min": 0, "max": 1024}), - } - } - - RETURN_TYPES = ("IMAGE", "MASK") - RETURN_NAMES = ("image", "mask") - FUNCTION = "draw" - CATEGORY = "ygo_tools" - - def hex_to_rgba(self, hex_color): - """将HEX颜色值转换为RGBA。支持#RGB,#RGBA,#RRGGBB,#RRGGBBAA格式""" - hex_color = hex_color.lstrip('#') - - if len(hex_color) == 3: # RGB格式 - r = int(hex_color[0] + hex_color[0], 16) / 255.0 - g = int(hex_color[1] + hex_color[1], 16) / 255.0 - b = int(hex_color[2] + hex_color[2], 16) / 255.0 - a = 1.0 - elif len(hex_color) == 4: # RGBA格式 - r = int(hex_color[0] + hex_color[0], 16) / 255.0 - g = int(hex_color[1] + hex_color[1], 16) / 255.0 - b = int(hex_color[2] + hex_color[2], 16) / 255.0 - a = int(hex_color[3] + hex_color[3], 16) / 255.0 - elif len(hex_color) == 6: # RRGGBB格式 - r = int(hex_color[0:2], 16) / 255.0 - g = int(hex_color[2:4], 16) / 255.0 - b = int(hex_color[4:6], 16) / 255.0 - a = 1.0 - elif len(hex_color) == 8: # RRGGBBAA格式 - r = int(hex_color[0:2], 16) / 255.0 - g = int(hex_color[2:4], 16) / 255.0 - b = int(hex_color[4:6], 16) / 255.0 - a = int(hex_color[6:8], 16) / 255.0 - else: - # 无效格式,使用白色 - r, g, b, a = 1.0, 1.0, 1.0, 1.0 - - return r, g, b, a - - def draw(self, length, width, color, angle, padding): - # 解析颜色 - r, g, b, a = self.hex_to_rgba(color) - - # 通过角度计算线条的实际尺寸和方向 - angle_rad = math.radians(angle) - - # 计算旋转后线条的边界框大小 - diagonal = math.sqrt(length**2 + width**2) - - # 考虑旋转后的边界,确保线条完全可见 - canvas_size = int(diagonal) + 2 * padding - - # 创建一个正方形画布 - image = torch.zeros((1, canvas_size, canvas_size, 4), dtype=torch.float32) - mask = torch.zeros((1, canvas_size, canvas_size), dtype=torch.float32) - - # 计算中心点 - center_x = canvas_size // 2 - center_y = canvas_size // 2 - - # 计算旋转后的线条坐标 - for y in range(canvas_size): - for x in range(canvas_size): - # 相对于中心点的坐标 - rel_x = x - center_x - rel_y = y - center_y - - # 将坐标旋转到线条方向 - rotated_x = rel_x * math.cos(-angle_rad) - rel_y * math.sin(-angle_rad) - rotated_y = rel_x * math.sin(-angle_rad) + rel_y * math.cos(-angle_rad) - - # 检查点是否在线条区域内 - if abs(rotated_y) <= width/2 and -length/2 <= rotated_x <= length/2: - # 设置点的颜色 - image[0, y, x, 0] = r - image[0, y, x, 1] = g - image[0, y, x, 2] = b - image[0, y, x, 3] = a - - # 更新掩码 - mask[0, y, x] = a - - return (image, mask) \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..c5fdd66 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,3 @@ +opencv-python>=4.5.0 +torch>=1.7.0 +numpy>=1.19.0 \ No newline at end of file diff --git a/tile_image.py b/tile_image.py deleted file mode 100644 index e44d5c0..0000000 --- a/tile_image.py +++ /dev/null @@ -1,55 +0,0 @@ -import torch -import numpy as np - -class TileImage: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "image": ("IMAGE",), - "mask": ("MASK",), - "num_tiles": ("INT", {"default": 2, "min": 1, "max": 12}), - "spacing_factor": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 2.0, "step": 0.001}), - "align": (["right", "left"], {"default": "right"}), - } - } - - RETURN_TYPES = ("IMAGE", "MASK") - RETURN_NAMES = ("image", "mask") - FUNCTION = "tile" - CATEGORY = "ygo_tools" - - def tile(self, image, mask, num_tiles, spacing_factor, align): - # 获取图像尺寸 - b, h, w, c = image.shape - - # 根据分辨率和因子计算间距 - spacing = int(w * spacing_factor) - - # 固定画布大小为12个图像平铺 - fixed_tiles = 12 - new_width = w * fixed_tiles + spacing * (fixed_tiles - 1) - new_image = torch.zeros((b, h, new_width, c), dtype=image.dtype) - - # 创建新的mask画布,填充为1(完全不透明) - new_mask = torch.ones((b, h, new_width), dtype=mask.dtype) - - # 限制实际平铺次数不超过12 - num_tiles = min(num_tiles, fixed_tiles) - - # 复制图像和mask并添加间距 - for i in range(num_tiles): - if align == "right": - # 从右往左排列 - tile_index = fixed_tiles - i - 1 - else: - # 从左往右排列 - tile_index = i - - start_x = tile_index * (w + spacing) - # 直接复制原始图像,保持透明通道不变 - new_image[:, :, start_x:start_x + w, :] = image - # 复制mask,间隔部分保持透明(0值) - new_mask[:, :, start_x:start_x + w] = mask - - return (new_image, new_mask) \ No newline at end of file