init code
This commit is contained in:
stormcenter
@@ -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
|
||||
}
|
||||
@@ -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.
|
||||
|
||||
@@ -0,0 +1,7 @@
|
||||
from .GridImageSplitter import GridImageSplitter, NODE_CLASS_MAPPINGS
|
||||
|
||||
__all__ = ['GridImageSplitter']
|
||||
|
||||
NODE_CLASS_MAPPINGS.update({
|
||||
"GridImageSplitter": GridImageSplitter
|
||||
})
|
||||
+217
@@ -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
|
||||
}
|
||||
@@ -0,0 +1,4 @@
|
||||
numpy>=1.21.0
|
||||
opencv-python>=4.5.0
|
||||
pillow>=8.0.0
|
||||
torch>=1.7.0
|
||||
Reference in New Issue
Block a user