Class MaskMerge

mask_split.py:140–260

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)