Class MaskSplit

mask_split.py:7–136

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]