diff --git a/comfy_extras/nodes_mask_conversion.py b/comfy_extras/nodes_mask_conversion.py new file mode 100644 index 00000000..04dcbd0d --- /dev/null +++ b/comfy_extras/nodes_mask_conversion.py @@ -0,0 +1,54 @@ +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image + +import comfy.utils + +class ImageToMask: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "channel": (["red", "green", "blue"],), + } + } + + CATEGORY = "image" + + RETURN_TYPES = ("MASK",) + FUNCTION = "image_to_mask" + + def image_to_mask(self, image, channel): + channels = ["red", "green", "blue"] + mask = image[0, :, :, channels.index(channel)] + return (mask,) + +class MaskToImage: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "mask": ("MASK",), + } + } + + CATEGORY = "image" + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "mask_to_image" + + def mask_to_image(self, mask): + result = mask[None, :, :, None].expand(-1, -1, -1, 3) + return (result,) + +NODE_CLASS_MAPPINGS = { + "ImageToMask": ImageToMask, + "MaskToImage": MaskToImage, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ImageToMask": "Convert Image to Mask", + "MaskToImage": "Convert Mask to Image", +} diff --git a/nodes.py b/nodes.py index 6468ac6b..aff03dd4 100644 --- a/nodes.py +++ b/nodes.py @@ -1193,3 +1193,4 @@ def init_custom_nodes(): load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_upscale_model.py")) load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_post_processing.py")) load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_mask.py")) + load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_mask_conversion.py"))