210 lines
8.3 KiB
Python
210 lines
8.3 KiB
Python
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}),
|
|
"row_split_method": (["uniform", "edge_detection"],),
|
|
"col_split_method": (["uniform", "edge_detection"],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "IMAGE")
|
|
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]
|
|
return [i * size // (num_splits + 1) for i in range(1, num_splits + 1)]
|
|
else: # edge_detection
|
|
gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
|
|
edges = cv2.Canny(gray, 50, 150)
|
|
edge_density = np.sum(edges, axis=0) if is_vertical else np.sum(edges, axis=1)
|
|
|
|
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)
|
|
return split_positions
|
|
|
|
def split_image(self, image, rows, cols, row_split_method, col_split_method):
|
|
print(f"Input image shape: {image.shape}")
|
|
print(f"Input image dtype: {image.dtype}")
|
|
|
|
if len(image.shape) == 3:
|
|
image = image.unsqueeze(0)
|
|
|
|
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 = 0
|
|
max_width = 0
|
|
temp_splits = []
|
|
|
|
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_np = img_np[top:bottom, left:right]
|
|
trimmed_cell = self.trim_white_borders(cell_np)
|
|
|
|
max_height = max(max_height, trimmed_cell.shape[0])
|
|
max_width = max(max_width, trimmed_cell.shape[1])
|
|
|
|
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)
|
|
|
|
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
|
|
}
|