From 21fc99760cf8ca56fc0e1f6ddbb963cd6bbde28b Mon Sep 17 00:00:00 2001 From: stormcenter Date: Sun, 3 Nov 2024 16:34:19 +0800 Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E7=99=BD=E8=89=B2=E8=BE=B9?= =?UTF-8?q?=E7=BC=98=E5=A4=84=E7=90=86?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- GridImageSplitter.py | 136 +++++++++++++++++++++++++++++++++++++------ 1 file changed, 119 insertions(+), 17 deletions(-) diff --git a/GridImageSplitter.py b/GridImageSplitter.py index f9fc5ca..373a850 100644 --- a/GridImageSplitter.py +++ b/GridImageSplitter.py @@ -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 }