diff --git a/GridImageSplitter.py b/GridImageSplitter.py new file mode 100644 index 0000000..5998a9c --- /dev/null +++ b/GridImageSplitter.py @@ -0,0 +1,108 @@ +import torch +import numpy as np +import cv2 + +class GridImageSplitter: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "rows": ("INT", {"default": 2, "min": 1, "max": 10}), + "cols": ("INT", {"default": 3, "min": 1, "max": 10}), + }, + } + + RETURN_TYPES = ("IMAGE", "IMAGE") + FUNCTION = "split_image" + CATEGORY = "image/processing" + + def find_split_positions(self, image, num_splits, is_vertical): + if is_vertical: + # 对于列，保持原有的边缘检测方法 + gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) + edges = cv2.Canny(gray, 50, 150) + edge_density = np.sum(edges, axis=0) + + window_size = len(edge_density) // (num_splits + 1) // 2 + smoothed_density = np.convolve(edge_density, np.ones(window_size)/window_size, mode='same') + + split_positions = [] + for i in range(1, num_splits + 1): + start = i * len(smoothed_density) // (num_splits + 1) - window_size + end = i * len(smoothed_density) // (num_splits + 1) + window_size + split = start + np.argmin(smoothed_density[start:end]) + split_positions.append(split) + else: + # 对于行，使用均匀分割 + height = image.shape[0] + split_positions = [i * height // (num_splits + 1) for i in range(1, num_splits + 1)] + + return split_positions + + def split_image(self, image, rows, cols): + print(f"Input image shape: {image.shape}") + print(f"Input image dtype: {image.dtype}") + + # 确保输入图像是 4D tensor (B, H, W, C) + if len(image.shape) == 3: + image = image.unsqueeze(0) + + # 将图像从 PyTorch tensor 转换为 NumPy 数组 + img_np = (image[0].cpu().numpy() * 255).astype(np.uint8) + + height, width = img_np.shape[:2] + print(f"Original image size: {width}x{height}") + + # 找到分割位置 + vertical_splits = self.find_split_positions(img_np, cols - 1, True) + horizontal_splits = self.find_split_positions(img_np, rows - 1, False) + + # 创建预览图 + preview_img = image.clone() + + # 在预览图上画线 + green_line = torch.tensor([0.0, 1.0, 0.0]).view(1, 1, 1, 3) + for x in vertical_splits: + preview_img[:, :, x:x+2, :] = green_line + for y in horizontal_splits: + preview_img[:, y:y+2, :, :] = green_line + + # 分割图片 + split_images = [] + h_splits = [0] + horizontal_splits + [height] + v_splits = [0] + vertical_splits + [width] + + # 计算最大的分割尺寸 + max_height = max([h_splits[i+1] - h_splits[i] for i in range(len(h_splits) - 1)]) + max_width = max([v_splits[i+1] - v_splits[i] for i in range(len(v_splits) - 1)]) + + for i in range(len(h_splits) - 1): + for j in range(len(v_splits) - 1): + top = h_splits[i] + bottom = h_splits[i+1] + left = v_splits[j] + right = v_splits[j+1] + + cell = image[:, top:bottom, left:right, :] + + # 填充到最大尺寸 + pad_bottom = max_height - (bottom - top) + pad_right = max_width - (right - left) + cell = torch.nn.functional.pad(cell, (0, 0, 0, pad_right, 0, pad_bottom), mode='constant') + + split_images.append(cell) + + split_images = torch.cat(split_images, dim=0) # 将所有分割图像合并为一个 tensor + + print(f"Preview tensor shape: {preview_img.shape}") + print(f"Preview tensor dtype: {preview_img.dtype}") + print(f"Split images tensor shape: {split_images.shape}") + print(f"Split images tensor dtype: {split_images.dtype}") + + return (preview_img, split_images) + +# 这个字典用于注册节点 +NODE_CLASS_MAPPINGS = { + "GridImageSplitter": GridImageSplitter +} diff --git a/README.md b/README.md index 8d81c05..f539105 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,53 @@ # ComfyUI-AutoSplitGridImage -ComfyUI-AutoSplitGridImage: A custom node for ComfyUI that intelligently splits images into grids, combining edge detection for columns and uniform division for rows. + +ComfyUI-AutoSplitGridImage is a custom node for ComfyUI that provides intelligent image splitting functionality. It combines edge detection for column splits and uniform division for row splits, offering a balanced approach to grid-based image segmentation. + +## Features + +- Intelligent column splitting using edge detection +- Uniform row splitting for consistent horizontal divisions +- Customizable number of rows and columns +- Preview image with grid lines +- Outputs both the preview image and individual grid cells + +## Installation + +1. Clone this repository into your ComfyUI `custom_nodes` directory: + ``` + cd /path/to/ComfyUI/custom_nodes + git clone https://github.com/yourusername/ComfyUI-AutoSplitGridImage.git + ``` +2. Restart ComfyUI or reload custom nodes. + +## Usage + +1. In the ComfyUI interface, find the "GridImageSplitter" node under the "image/processing" category. +2. Connect an image output to the "image" input of the GridImageSplitter node. +3. Set the desired number of rows and columns. +4. The node will output two images: + - A preview image showing the grid lines + - A tensor containing all the split image cells + +## Parameters + +- `image`: Input image to be split +- `rows`: Number of rows to split the image into (default: 2, range: 1-10) +- `cols`: Number of columns to split the image into (default: 3, range: 1-10) + +## How It Works + +- Column Splitting: Uses edge detection to find natural splitting points in the image. +- Row Splitting: Applies uniform division for consistent horizontal splits. + +## Contributing + +Contributions are welcome! Please feel free to submit a Pull Request. + +## License + +This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details. + +## Acknowledgements + +- This project is designed to work with [ComfyUI](https://github.com/comfyanonymous/ComfyUI). +- Edge detection is implemented using OpenCV. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..19f4fbd --- /dev/null +++ b/__init__.py @@ -0,0 +1,7 @@ +from .GridImageSplitter import GridImageSplitter, NODE_CLASS_MAPPINGS + +__all__ = ['GridImageSplitter'] + +NODE_CLASS_MAPPINGS.update({ + "GridImageSplitter": GridImageSplitter +}) \ No newline at end of file diff --git a/example.json b/example.json new file mode 100644 index 0000000..6df3e86 --- /dev/null +++ b/example.json @@ -0,0 +1,217 @@ +{ + "last_node_id": 7, + "last_link_id": 3, + "nodes": [ + { + "id": 4, + "type": "LoadImage", + "pos": { + "0": 970, + "1": 456, + "2": 0, + "3": 0, + "4": 0, + "5": 0, + "6": 0, + "7": 0, + "8": 0, + "9": 0 + }, + "size": { + "0": 315, + "1": 314 + }, + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 1 + ], + "slot_index": 0, + "shape": 3 + }, + { + "name": "MASK", + "type": "MASK", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "b6eea9a635e147bdbec1d2ebfa22c6e0 (1).png", + "image" + ] + }, + { + "id": 5, + "type": "PreviewImage", + "pos": { + "0": 1920, + "1": 199, + "2": 0, + "3": 0, + "4": 0, + "5": 0, + "6": 0, + "7": 0, + "8": 0, + "9": 0 + }, + "size": { + "0": 210, + "1": 246 + }, + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 2 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "PreviewImage" + } + }, + { + "id": 3, + "type": "GridImageSplitter", + "pos": { + "0": 1371, + "1": 491, + "2": 0, + "3": 0, + "4": 0, + "5": 0, + "6": 0, + "7": 0, + "8": 0, + "9": 0 + }, + "size": { + "0": 315, + "1": 150 + }, + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 1 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 2 + ], + "slot_index": 0, + "shape": 3 + }, + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 3 + ], + "slot_index": 1, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "GridImageSplitter" + }, + "widgets_values": [ + 2, + 3 + ] + }, + { + "id": 6, + "type": "PreviewImage", + "pos": { + "0": 1923, + "1": 715, + "2": 0, + "3": 0, + "4": 0, + "5": 0, + "6": 0, + "7": 0, + "8": 0, + "9": 0 + }, + "size": { + "0": 210, + "1": 246 + }, + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 3 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "PreviewImage" + } + } + ], + "links": [ + [ + 1, + 4, + 0, + 3, + 0, + "IMAGE" + ], + [ + 2, + 3, + 0, + 5, + 0, + "IMAGE" + ], + [ + 3, + 3, + 1, + 6, + 0, + "IMAGE" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.7247295000000008, + "offset": [ + -224.91286373867808, + -64.33740337427317 + ] + } + }, + "version": 0.4 +} \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..428aa8e --- /dev/null +++ b/requirements.txt @@ -0,0 +1,4 @@ +numpy>=1.21.0 +opencv-python>=4.5.0 +pillow>=8.0.0 +torch>=1.7.0 \ No newline at end of file