diff --git a/README.md b/README.md
index 59a0d1f..cf7467f 100644
--- a/README.md
+++ b/README.md
@@ -1,13 +1,18 @@
# comfyui-easyapi-nodes
-针对api接口开发补充的一些自定义节点和功能
+针对api接口开发补充的一些自定义节点和功能。
+
+转成base64的节点都是输出节点,websocket消息中会包含base64Images和base64Type属性(具体格式请查看ImageNode.py中的ImageToBase64Advanced类源代码,或者自己搭建简单流程运行在浏览器开发者工具-->网络中查看)
## 节点
-| 名称 | 说明 |
-|-----|------------------------------------------------------------------------------------------------------------------------------------------------------------|
-| Base64ToImage | 把图片base64字符串转成图片 |
-| ImageToBase64 | 把图片转成base64字符串 |
-| LoadImageToBase64 | 加载本地图片转成base64字符串 |
-| SamAutoMaskSEGS | 得到图片所有语义分割的coco或uncompress_rle格式。
配合ComfyUI-Impact-Pack的SAMLoader或comfyui_segment_anything的SAMModelLoader。
但是如果使用hq模型,必须使用comfyui_segment_anything |
+| 名称 | 说明 |
+|-----------------------|------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| Base64ToImage | 把图片base64字符串转成图片 |
+| ImageToBase64Advanced | 把图片转成base64字符串, 可以选择图片类型(image, mask) ,方便接口调用判断 |
+| ImageToBase64 | 把图片转成base64字符串(imageType=["image"]) |
+| MaskToBase64Image | 把遮罩转成对应图片的base64字符串(imageType=["mask"]) |
+| MaskImageToBase64 | 把遮罩图片转成base64字符串(imageType=["mask"]) |
+| LoadImageToBase64 | 加载本地图片转成base64字符串 |
+| SamAutoMaskSEGS | 得到图片所有语义分割的coco或uncompress_rle格式。
配合ComfyUI-Impact-Pack的SAMLoader或comfyui_segment_anything的SAMModelLoader。
但是如果使用hq模型,必须使用comfyui_segment_anything |
## 功能
- 扩展Save(Api Format)菜单。
diff --git a/easyapi/ImageNode.py b/easyapi/ImageNode.py
index 8f732c5..9374b0f 100644
--- a/easyapi/ImageNode.py
+++ b/easyapi/ImageNode.py
@@ -10,6 +10,7 @@ import json
from json import JSONEncoder, JSONDecoder
from easyapi.util import tensor_to_pil
+
class Base64ToImage:
@classmethod
def INPUT_TYPES(self):
@@ -50,11 +51,15 @@ class Base64ToImage:
# return (torch.stack(images, dim=0)[None, ], mask.unsqueeze(0))
-class ImageToBase64:
+class ImageToBase64Advanced:
+ def __init__(self):
+ self.imageType = "image"
+
@classmethod
def INPUT_TYPES(self):
return {"required": {
"images": ("IMAGE",),
+ "imageType": (["image", "mask"], {"default": "image"}),
},
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
}
@@ -72,7 +77,10 @@ class ImageToBase64:
# INPUT_IS_LIST = False
# OUTPUT_IS_LIST = (False,False,)
- def convert(self, images, prompt=None, extra_pnginfo=None):
+ def convert(self, images, imageType=None, prompt=None, extra_pnginfo=None):
+ if imageType is None:
+ imageType = self.imageType
+
result = list()
for i in images:
img = tensor_to_pil(i)
@@ -99,11 +107,47 @@ class ImageToBase64:
result.append(encoded_image)
base64Images = JSONEncoder().encode(result)
# print(images)
- return {"ui": {"base64Images": result}, "result": (base64Images,)}
+ return {"ui": {"base64Images": result, "imageType": [imageType]}, "result": (base64Images,)}
+
+
+class ImageToBase64(ImageToBase64Advanced):
+ def __init__(self):
+ self.imageType = "image"
+
+ @classmethod
+ def INPUT_TYPES(self):
+ return {"required": {
+ "images": ("IMAGE",),
+ },
+ "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
+ }
+
+
+class MaskImageToBase64(ImageToBase64):
+ def __init__(self):
+ self.imageType = "mask"
+
+
+class MaskToBase64Image(MaskImageToBase64):
+ @classmethod
+ def INPUT_TYPES(s):
+ return {
+ "required": {
+ "mask": ("MASK",),
+ }
+ }
+
+ CATEGORY = "EasyApi/Image"
+
+ RETURN_TYPES = ("STRING",)
+ FUNCTION = "mask_to_base64image"
+
+ def mask_to_base64image(self, mask):
+ images = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
+ return super().convert(images)
class LoadImageToBase64(LoadImage):
-
RETURN_TYPES = ("STRING", "IMAGE", "MASK", )
RETURN_NAMES = ("base64Images", "IMAGE", "MASK", )
@@ -152,6 +196,9 @@ def base64_to_image(base64_string):
NODE_CLASS_MAPPINGS = {
"Base64ToImage": Base64ToImage,
"ImageToBase64": ImageToBase64,
+ "ImageToBase64Advanced": ImageToBase64Advanced,
+ "MaskToBase64Image": MaskToBase64Image,
+ "MaskImageToBase64": MaskImageToBase64,
"LoadImageToBase64": LoadImageToBase64,
}
@@ -159,5 +206,8 @@ NODE_CLASS_MAPPINGS = {
NODE_DISPLAY_NAME_MAPPINGS = {
"Base64ToImage": "Base64 To Image",
"ImageToBase64": "Image To Base64",
+ "ImageToBase64Advanced": "Image To Base64 (Advanced)",
+ "MaskToBase64Image": "Mask To Base64 Image",
+ "MaskImageToBase64": "Mask Image To Base64",
"LoadImageToBase64": "Load Image To Base64",
}