mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2025-01-11 02:15:17 +00:00
Merge branch 'image-to-mask' of https://github.com/missionfloyd/ComfyUI
# Conflicts: # nodes.py
This commit is contained in:
commit
e1db7a2038
54
comfy_extras/nodes_mask_conversion.py
Normal file
54
comfy_extras/nodes_mask_conversion.py
Normal file
@ -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",
|
||||
}
|
1
nodes.py
1
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"))
|
||||
|
Loading…
Reference in New Issue
Block a user