mask_split.py

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🐰",
}