296 lines
11 KiB
Python
296 lines
11 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 remove_external_borders(self, img_np):
|
|
"""
|
|
处理外部边缘，加强对黑色边框的检测
|
|
"""
|
|
if img_np.size == 0 or img_np is None:
|
|
return img_np
|
|
|
|
# 转换为多个色彩空间
|
|
gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
|
|
hsv = cv2.cvtColor(img_np, cv2.COLOR_RGB2HSV)
|
|
|
|
height, width = img_np.shape[:2]
|
|
|
|
# 最小裁剪量
|
|
min_trim = 18
|
|
# 检查范围扩大到30像素
|
|
check_width = 30
|
|
|
|
def is_black_region(region):
|
|
"""检查区域是否为黑色区域"""
|
|
if len(region.shape) == 3:
|
|
region_gray = cv2.cvtColor(region, cv2.COLOR_RGB2GRAY)
|
|
else:
|
|
region_gray = region
|
|
|
|
# 计算暗色像素的比例
|
|
dark_ratio = np.mean(region_gray < 40)
|
|
# 如果超过60%的像素是暗色的，认为是黑边
|
|
return dark_ratio > 0.6
|
|
|
|
def is_white_region(region):
|
|
"""检查区域是否为白色区域"""
|
|
if len(region.shape) == 3:
|
|
region_hsv = cv2.cvtColor(region, cv2.COLOR_RGB2HSV)
|
|
sat_mean = np.mean(region_hsv[:,:,1])
|
|
val_mean = np.mean(region_hsv[:,:,2])
|
|
return sat_mean < 30 and val_mean > 225
|
|
return False
|
|
|
|
def find_border(gray_img, is_vertical=True, from_start=True):
|
|
"""查找边界"""
|
|
if is_vertical:
|
|
total_size = width
|
|
chunk_size = 5 # 每次检查5个像素
|
|
else:
|
|
total_size = height
|
|
chunk_size = 5
|
|
|
|
if from_start:
|
|
range_iter = range(0, total_size-chunk_size, chunk_size)
|
|
else:
|
|
range_iter = range(total_size-chunk_size, 0, -chunk_size)
|
|
|
|
for i in range_iter:
|
|
if is_vertical:
|
|
chunk = img_np[:, i:i+chunk_size] if from_start else img_np[:, i-chunk_size:i]
|
|
else:
|
|
chunk = img_np[i:i+chunk_size, :] if from_start else img_np[i-chunk_size:i, :]
|
|
|
|
# 分别检查黑边和白边
|
|
if not (is_black_region(chunk) or is_white_region(chunk)):
|
|
return i if from_start else i
|
|
|
|
return min_trim if from_start else total_size - min_trim
|
|
|
|
# 检测左边界
|
|
left = find_border(gray, is_vertical=True, from_start=True)
|
|
|
|
# 检测右边界
|
|
right = find_border(gray, is_vertical=True, from_start=False)
|
|
|
|
# 检测上边界
|
|
top = find_border(gray, is_vertical=False, from_start=True)
|
|
|
|
# 检测下边界
|
|
bottom = find_border(gray, is_vertical=False, from_start=False)
|
|
|
|
# 强制应用最小裁剪
|
|
# 检查左侧边缘
|
|
left_region = gray[:, :check_width]
|
|
if np.mean(left_region < 40) > 0.3: # 如果有超过30%的暗色像素
|
|
left = max(left, min_trim)
|
|
|
|
# 检查右侧边缘
|
|
right_region = gray[:, -check_width:]
|
|
if np.mean(right_region < 40) > 0.3: # 如果有超过30%的暗色像素
|
|
right = min(right, width - min_trim)
|
|
|
|
# 检查上边缘
|
|
top_region = gray[:check_width, :]
|
|
if np.mean(top_region < 40) > 0.3:
|
|
top = max(top, min_trim)
|
|
|
|
# 检查下边缘
|
|
bottom_region = gray[-check_width:, :]
|
|
if np.mean(bottom_region < 40) > 0.3:
|
|
bottom = min(bottom, height - min_trim)
|
|
|
|
# 确保裁剪合理
|
|
if (right - left) < width * 0.5 or (bottom - top) < height * 0.5:
|
|
return img_np
|
|
|
|
# 应用裁剪
|
|
cropped = img_np[top:bottom, left:right]
|
|
|
|
# 进行二次检查，确保没有遗漏的黑边
|
|
if cropped.shape[1] > 2 * min_trim:
|
|
right_edge = cv2.cvtColor(cropped[:, -min_trim:], cv2.COLOR_RGB2GRAY)
|
|
if np.mean(right_edge < 40) > 0.3:
|
|
cropped = cropped[:, :-min_trim]
|
|
|
|
return cropped
|
|
|
|
def detect_split_borders(self, img_strip, is_vertical=True):
|
|
"""
|
|
检测分割线区域的边框
|
|
"""
|
|
hsv = cv2.cvtColor(img_strip, cv2.COLOR_RGB2HSV)
|
|
lab = cv2.cvtColor(img_strip, cv2.COLOR_RGB2LAB)
|
|
|
|
sat = hsv[:, :, 1]
|
|
val = hsv[:, :, 2]
|
|
l_channel = lab[:, :, 0]
|
|
|
|
# 检测白色和黑色区域
|
|
is_border = ((sat < 10) & (val > 248)) | (l_channel < 30)
|
|
|
|
if is_vertical:
|
|
border_ratios = np.mean(is_border, axis=1)
|
|
indices = np.where(border_ratios < 0.95)[0]
|
|
else:
|
|
border_ratios = np.mean(is_border, axis=0)
|
|
indices = np.where(border_ratios < 0.95)[0]
|
|
|
|
if len(indices) == 0:
|
|
return 0, img_strip.shape[1] if is_vertical else img_strip.shape[0]
|
|
|
|
return indices[0], indices[-1]
|
|
|
|
def adjust_split_line(self, img_np, split_pos, is_vertical=True, margin=15):
|
|
"""
|
|
调整分割线附近的边界
|
|
"""
|
|
height, width = img_np.shape[:2]
|
|
|
|
if is_vertical:
|
|
left_bound = max(0, split_pos - margin)
|
|
right_bound = min(width, split_pos + margin)
|
|
strip = img_np[:, left_bound:right_bound]
|
|
start, end = self.detect_split_borders(strip, False)
|
|
|
|
start = max(0, start - 5)
|
|
end = min(strip.shape[1], end + 5)
|
|
|
|
return left_bound + start, left_bound + end
|
|
else:
|
|
top_bound = max(0, split_pos - margin)
|
|
bottom_bound = min(height, split_pos + margin)
|
|
strip = img_np[top_bound:bottom_bound, :]
|
|
start, end = self.detect_split_borders(strip, True)
|
|
|
|
start = max(0, start - 5)
|
|
end = min(strip.shape[0], end + 5)
|
|
|
|
return top_bound + start, top_bound + end
|
|
|
|
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:
|
|
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):
|
|
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]
|
|
|
|
# 获取分割位置
|
|
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
|
|
|
|
# 调整分割位置
|
|
adjusted_v_splits = []
|
|
for split in vertical_splits:
|
|
left, right = self.adjust_split_line(img_np, split, True)
|
|
adjusted_v_splits.extend([left, right])
|
|
|
|
adjusted_h_splits = []
|
|
for split in horizontal_splits:
|
|
top, bottom = self.adjust_split_line(img_np, split, False)
|
|
adjusted_h_splits.extend([top, bottom])
|
|
|
|
# 获取分割边界
|
|
h_splits = [0] + sorted(adjusted_h_splits) + [height]
|
|
v_splits = [0] + sorted(adjusted_v_splits) + [width]
|
|
|
|
# 处理所有分割区域
|
|
split_images = []
|
|
max_h = 0
|
|
max_w = 0
|
|
|
|
# 第一次遍历: 找出所有裁剪后图片的最大宽度和高度
|
|
temp_splits = []
|
|
for i in range(0, len(h_splits)-1, 2):
|
|
for j in range(0, len(v_splits)-1, 2):
|
|
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.remove_external_borders(cell_np)
|
|
temp_splits.append(trimmed_cell)
|
|
|
|
h, w = trimmed_cell.shape[:2]
|
|
max_h = max(max_h, h)
|
|
max_w = max(max_w, w)
|
|
|
|
# 第二次遍历: 将所有图片调整为相同尺寸，保持原始比例
|
|
for cell_np in temp_splits:
|
|
h, w = cell_np.shape[:2]
|
|
# 计算缩放比例
|
|
scale = min(max_h/h, max_w/w)
|
|
new_h = int(h * scale)
|
|
new_w = int(w * scale)
|
|
|
|
# 居中放置
|
|
resized = cv2.resize(cell_np, (new_w, new_h), interpolation=cv2.INTER_LANCZOS4)
|
|
canvas = np.zeros((max_h, max_w, 3), dtype=np.uint8)
|
|
y_offset = (max_h - new_h) // 2
|
|
x_offset = (max_w - new_w) // 2
|
|
canvas[y_offset:y_offset+new_h, x_offset:x_offset+new_w] = resized
|
|
|
|
# 转换为tensor
|
|
cell_tensor = torch.from_numpy(canvas).float() / 255.0
|
|
cell_tensor = cell_tensor.unsqueeze(0)
|
|
split_images.append(cell_tensor)
|
|
|
|
stacked_images = torch.cat(split_images, dim=0)
|
|
|
|
if stacked_images.shape[-1] != 3:
|
|
stacked_images = stacked_images.permute(0, 2, 3, 1)
|
|
|
|
print(f"Final stacked shape: {stacked_images.shape}")
|
|
|
|
return (preview_img, stacked_images)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"GridImageSplitter": GridImageSplitter
|
|
}
|