nodes/stormcenter-ComfyUI-AutoSplitGridImage
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

增加白色边缘处理

Browse Source
This commit is contained in:
stormcenter
2024-11-03 16:34:19 +08:00
parent 6d6a2511d9
commit 21fc99760c
1 changed files with 119 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
GridImageSplitter.py
+119 -17
View File
@@ -19,6 +19,106 @@ class GridImageSplitter:
FUNCTION = "split_image"
CATEGORY = "image/processing"
def trim_white_borders(self, img_np):
"""
增强的白边裁剪算法，特别优化了对天空等浅色区域的处理
"""
if len(img_np.shape) == 3:
# 转换为多个颜色空间以更好地检测边界
hsv = cv2.cvtColor(img_np, cv2.COLOR_RGB2HSV)
lab = cv2.cvtColor(img_np, cv2.COLOR_RGB2LAB)
# 提取各个通道
v_channel = hsv[:,:,2]
l_channel = lab[:,:,0]
# 计算梯度
gradient_v = cv2.Sobel(v_channel, cv2.CV_64F, 1, 1, ksize=3)
gradient_l = cv2.Sobel(l_channel, cv2.CV_64F, 1, 1, ksize=3)
# 综合多个指标
gradient_magnitude = np.sqrt(gradient_v**2 + gradient_l**2)
# 创建掩码
mask = np.zeros_like(v_channel)
# 基于梯度和亮度的自适应阈值
gradient_threshold = np.percentile(gradient_magnitude, 95) * 0.1
brightness_threshold = np.percentile(v_channel, 95) * 0.95
# 更新掩码
mask[gradient_magnitude > gradient_threshold] = 255
mask[v_channel < brightness_threshold] = 255
# 应用形态学操作
kernel = np.ones((3,3), np.uint8)
mask = cv2.dilate(mask, kernel, iterations=1)
mask = cv2.erode(mask, kernel, iterations=1)
else:
# 如果是灰度图像
gradient = cv2.Sobel(img_np, cv2.CV_64F, 1, 1, ksize=3)
mask = (gradient > np.percentile(gradient, 95) * 0.1).astype(np.uint8) * 255
# 找到非边界区域
coords = cv2.findNonZero(mask)
if coords is None:
return img_np
x, y, w, h = cv2.boundingRect(coords)
# 智能边界调整
def adjust_boundary(start, end, size, is_start=True):
if is_start:
while start < end and start < size - 1:
if np.any(mask[start:start+1, :] > 0):
break
start += 1
return max(0, start - 1)
else:
while end > start and end > 0:
if np.any(mask[end-1:end, :] > 0):
break
end -= 1
return min(size, end + 1)
# 调整边界
y = adjust_boundary(y, y+h, img_np.shape[0], True)
h = adjust_boundary(y, y+h, img_np.shape[0], False) - y
x = adjust_boundary(x, x+w, img_np.shape[1], True)
w = adjust_boundary(x, x+w, img_np.shape[1], False) - x
# 裁剪图像
cropped = img_np[y:y+h, x:x+w]
# 边缘检查和优化
if cropped.shape[0] > 10 and cropped.shape[1] > 10:
edge_width = 3
edges = {
'top': cropped[:edge_width, :],
'bottom': cropped[-edge_width:, :],
'left': cropped[:, :edge_width],
'right': cropped[:, -edge_width:]
}
for edge_name, edge in edges.items():
if len(edge.shape) == 3:
edge_gradient = cv2.Sobel(cv2.cvtColor(edge, cv2.COLOR_RGB2LAB)[:,:,0],
cv2.CV_64F, 1, 1, ksize=3)
else:
edge_gradient = cv2.Sobel(edge, cv2.CV_64F, 1, 1, ksize=3)
if np.mean(np.abs(edge_gradient)) < gradient_threshold:
if edge_name == 'top' and y + edge_width < y + h:
cropped = cropped[edge_width:, :]
elif edge_name == 'bottom' and h - edge_width > 0:
cropped = cropped[:-edge_width, :]
elif edge_name == 'left' and x + edge_width < x + w:
cropped = cropped[:, edge_width:]
elif edge_name == 'right' and w - edge_width > 0:
cropped = cropped[:, :-edge_width]
return cropped
def find_split_positions(self, image, num_splits, is_vertical, split_method):
if split_method == "uniform":
size = image.shape[1] if is_vertical else image.shape[0]
@@ -43,38 +143,32 @@ class GridImageSplitter:
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, col_split_method)
horizontal_splits = self.find_split_positions(img_np, rows - 1, False, row_split_method)
# 创建预览图
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)])
max_height = 0
max_width = 0
temp_splits = []
for i in range(len(h_splits) - 1):
for j in range(len(v_splits) - 1):
@@ -83,16 +177,25 @@ class GridImageSplitter:
left = v_splits[j]
right = v_splits[j+1]
cell = image[:, top:bottom, left:right, :]
cell_np = img_np[top:bottom, left:right]
trimmed_cell = self.trim_white_borders(cell_np)
# 填充到最大尺寸
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')
max_height = max(max_height, trimmed_cell.shape[0])
max_width = max(max_width, trimmed_cell.shape[1])
split_images.append(cell)
temp_splits.append(trimmed_cell)
for trimmed_cell in temp_splits:
cell_tensor = torch.from_numpy(trimmed_cell).float() / 255.0
cell_tensor = cell_tensor.unsqueeze(0)
pad_bottom = max_height - trimmed_cell.shape[0]
pad_right = max_width - trimmed_cell.shape[1]
cell_tensor = torch.nn.functional.pad(cell_tensor, (0, 0, 0, pad_right, 0, pad_bottom), mode='constant')
split_images.append(cell_tensor)
split_images = torch.cat(split_images, dim=0) # 将所有分割图像合并为一个 tensor
split_images = torch.cat(split_images, dim=0)
print(f"Preview tensor shape: {preview_img.shape}")
print(f"Preview tensor dtype: {preview_img.dtype}")
@@ -101,7 +204,6 @@ class GridImageSplitter:
return (preview_img, split_images)
# 这个字典用于注册节点
NODE_CLASS_MAPPINGS = {
"GridImageSplitter": GridImageSplitter
}