273 lines
9.0 KiB
Python
273 lines
9.0 KiB
Python
import torch
|
|
import torch.nn.functional as F
|
|
import cv2
|
|
import numpy as np
|
|
|
|
|
|
class MaskSplit:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"mask": ("MASK",),
|
|
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE","MASK")
|
|
RETURN_NAMES = ("segmented_images","segmented_masks")
|
|
FUNCTION = "segment_mask"
|
|
|
|
CATEGORY = "CyberEveLoop🐰"
|
|
|
|
def find_top_left_point(self, mask_np):
|
|
"""找到mask中最左上角的点"""
|
|
# 找到所有非零点
|
|
y_coords, x_coords = np.nonzero(mask_np)
|
|
if len(x_coords) == 0:
|
|
return float('inf'), float('inf')
|
|
|
|
# 找到最小x值
|
|
min_x = np.min(x_coords)
|
|
# 在最小x值的点中找到最小y值
|
|
min_y = np.min(y_coords[x_coords == min_x])
|
|
|
|
return min_x, min_y
|
|
|
|
def segment_mask(self, mask, image):
|
|
"""使用OpenCV快速分割蒙版并处理图像"""
|
|
# 保存原始设备信息
|
|
device = mask.device if isinstance(mask, torch.Tensor) else torch.device('cpu')
|
|
|
|
# 确保mask是正确的形状并转换为numpy数组
|
|
if isinstance(mask, torch.Tensor):
|
|
if len(mask.shape) == 2:
|
|
mask = mask.unsqueeze(0)
|
|
mask_np = (mask[0] * 255).cpu().numpy().astype(np.uint8)
|
|
else:
|
|
mask_np = (mask * 255).astype(np.uint8)
|
|
|
|
# 使用OpenCV找到轮廓
|
|
contours, hierarchy = cv2.findContours(
|
|
mask_np,
|
|
cv2.RETR_TREE,
|
|
cv2.CHAIN_APPROX_SIMPLE
|
|
)
|
|
|
|
mask_info = [] # 用于排序的信息列表
|
|
|
|
if hierarchy is not None and len(contours) > 0:
|
|
hierarchy = hierarchy[0]
|
|
contour_masks = {}
|
|
|
|
# 创建每个轮廓的mask
|
|
for i, contour in enumerate(contours):
|
|
mask = np.zeros_like(mask_np)
|
|
cv2.drawContours(mask, [contour], -1, 255, -1)
|
|
contour_masks[i] = mask
|
|
|
|
# 处理每个轮廓
|
|
processed_indices = set()
|
|
|
|
for i, (contour, h) in enumerate(zip(contours, hierarchy)):
|
|
if i in processed_indices:
|
|
continue
|
|
|
|
current_mask = contour_masks[i].copy()
|
|
child_idx = h[2]
|
|
|
|
if child_idx != -1:
|
|
while child_idx != -1:
|
|
current_mask = cv2.subtract(current_mask, contour_masks[child_idx])
|
|
processed_indices.add(child_idx)
|
|
child_idx = hierarchy[child_idx][0]
|
|
|
|
# 找到最左上角的点
|
|
min_x, min_y = self.find_top_left_point(current_mask)
|
|
|
|
# 转换为tensor
|
|
mask_tensor = torch.from_numpy(current_mask).float() / 255.0
|
|
mask_tensor = mask_tensor.unsqueeze(0)
|
|
mask_tensor = mask_tensor.to(device)
|
|
|
|
# 保存mask和排序信息
|
|
mask_info.append((mask_tensor, min_x, min_y))
|
|
processed_indices.add(i)
|
|
|
|
# 如果没有找到任何轮廓，使用原始mask
|
|
if not mask_info:
|
|
if isinstance(mask, torch.Tensor):
|
|
mask_info.append((mask, 0, 0))
|
|
else:
|
|
mask_tensor = torch.from_numpy(mask).float()
|
|
if len(mask_tensor.shape) == 2:
|
|
mask_tensor = mask_tensor.unsqueeze(0)
|
|
mask_tensor = mask_tensor.to(device)
|
|
mask_info.append((mask_tensor, 0, 0))
|
|
|
|
# 根据最左上角点排序
|
|
mask_info.sort(key=lambda x: (x[1], x[2]))
|
|
|
|
# 确保image是正确的形状
|
|
if len(image.shape) == 3:
|
|
image = image.unsqueeze(0)
|
|
|
|
# 处理masks和images
|
|
result_masks = None
|
|
result_images = None
|
|
|
|
for mask_tensor, _, _ in mask_info:
|
|
# 处理masks
|
|
if result_masks is None:
|
|
result_masks = mask_tensor
|
|
else:
|
|
result_masks = torch.cat([result_masks, mask_tensor], dim=0)
|
|
|
|
# 处理images
|
|
if result_images is None:
|
|
result_images = image.clone()
|
|
else:
|
|
result_images = torch.cat([result_images, image.clone()], dim=0)
|
|
|
|
return (result_images, result_masks)
|
|
|
|
|
|
|
|
class MaskMerge:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"original_image": ("IMAGE",),
|
|
},
|
|
"optional": {
|
|
"processed_images": ("IMAGE", {"forceInput": True}),
|
|
"masks": ("MASK", {"forceInput": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("merged_image",)
|
|
FUNCTION = "merge_masked_images"
|
|
CATEGORY = "CyberEveLoop🐰"
|
|
|
|
def standardize_input(self, image, processed_images=None, masks=None):
|
|
"""
|
|
标准化输入格式
|
|
- image: [H,W,C] -> [1,H,W,C]
|
|
- processed_images: [...] -> [B,H,W,C]
|
|
- masks: [...] -> [B,H,W]
|
|
"""
|
|
# 处理原始图像
|
|
if len(image.shape) == 3:
|
|
image = image.unsqueeze(0)
|
|
assert len(image.shape) == 4, f"Original image must be 4D [B,H,W,C], got shape {image.shape}"
|
|
|
|
# 处理processed_images
|
|
if processed_images is not None:
|
|
if isinstance(processed_images, list):
|
|
processed_images = torch.cat(processed_images, dim=0)
|
|
if len(processed_images.shape) == 3:
|
|
processed_images = processed_images.unsqueeze(0)
|
|
assert len(processed_images.shape) == 4, \
|
|
f"Processed images must be 4D [B,H,W,C], got shape {processed_images.shape}"
|
|
|
|
# 处理masks
|
|
if masks is not None:
|
|
if isinstance(masks, list):
|
|
masks = torch.cat(masks, dim=0)
|
|
if len(masks.shape) == 2:
|
|
masks = masks.unsqueeze(0)
|
|
assert len(masks.shape) == 3, f"Masks must be 3D [B,H,W], got shape {masks.shape}"
|
|
|
|
return image, processed_images, masks
|
|
|
|
def resize_tensor(self, x, size, mode='bilinear'):
|
|
"""调整tensor尺寸的辅助函数"""
|
|
# 确保输入是4D tensor [B,C,H,W]
|
|
orig_dim = x.dim()
|
|
if orig_dim == 3:
|
|
x = x.unsqueeze(0)
|
|
|
|
# 如果是图像 [B,H,W,C]，需要转换为 [B,C,H,W]
|
|
if x.shape[-1] in [1, 3, 4]:
|
|
x = x.permute(0, 3, 1, 2)
|
|
|
|
# 执行调整
|
|
x = F.interpolate(x, size=size, mode=mode, align_corners=False if mode in ['bilinear', 'bicubic'] else None)
|
|
|
|
# 转换回原始格式
|
|
if x.shape[1] in [1, 3, 4]:
|
|
x = x.permute(0, 2, 3, 1)
|
|
|
|
# 如果原始输入是3D，去掉batch维度
|
|
if orig_dim == 3:
|
|
x = x.squeeze(0)
|
|
|
|
return x
|
|
|
|
def merge_masked_images(self, original_image, processed_images=None, masks=None):
|
|
"""合并处理后的图像"""
|
|
# 确保输入有效
|
|
if processed_images is None or masks is None:
|
|
return (original_image,)
|
|
|
|
# 标准化输入
|
|
original_image, processed_images, masks = self.standardize_input(
|
|
original_image, processed_images, masks
|
|
)
|
|
|
|
# 创建结果图像的副本
|
|
result = original_image.clone()
|
|
|
|
# 获取目标尺寸
|
|
target_height = original_image.shape[1]
|
|
target_width = original_image.shape[2]
|
|
|
|
# 调整处理图像的尺寸（如果需要）
|
|
if processed_images.shape[1:3] != (target_height, target_width):
|
|
processed_images = self.resize_tensor(
|
|
processed_images,
|
|
(target_height, target_width),
|
|
mode='bilinear'
|
|
)
|
|
|
|
# 调整蒙版尺寸（如果需要）
|
|
if masks.shape[1:3] != (target_height, target_width):
|
|
masks = self.resize_tensor(
|
|
masks,
|
|
(target_height, target_width),
|
|
mode='bilinear'
|
|
)
|
|
|
|
# 扩展蒙版维度以匹配图像通道
|
|
masks = masks.unsqueeze(-1).expand(-1, -1, -1, 3)
|
|
|
|
# 批量处理所有图片
|
|
for i in range(processed_images.shape[0]):
|
|
current_image = processed_images[i:i+1]
|
|
current_mask = masks[i:i+1]
|
|
result = current_mask * current_image + (1 - current_mask) * result
|
|
|
|
assert len(result.shape) == 4, "Output must be 4D [B,H,W,C]"
|
|
return (result,)
|
|
|
|
|
|
Mask_CLASS_MAPPINGS = {
|
|
"CyberEve_MaskSegmentation": MaskSplit,
|
|
"CyberEve_MaskMerge": MaskMerge,
|
|
}
|
|
|
|
Mask_DISPLAY_NAME_MAPPINGS = {
|
|
"CyberEve_MaskSegmentation": "Mask Segmentation🐰",
|
|
"CyberEve_MaskMerge": "Mask Merge🐰",
|
|
}
|
|
|