返回消息区分mask和image

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格式。<br/>配合ComfyUI-Impact-Pack的SAMLoader或comfyui_segment_anything的SAMModelLoader。<br/>但是如果使用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格式。<br/>配合ComfyUI-Impact-Pack的SAMLoader或comfyui_segment_anything的SAMModelLoader。<br/>但是如果使用hq模型，必须使用comfyui_segment_anything |
 
 ## 功能
 - 扩展Save(Api Format)菜单。

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