diff --git a/README.md b/README.md index 089014e..a854bc2 100644 --- a/README.md +++ b/README.md @@ -45,13 +45,25 @@ | ShowBoolean | 显示布尔值(可指定消息中key值) | | ImageEqual | 图片是否相等（可用于通过判断遮罩图是否全黑来判定是否有遮罩） | | SDBaseVerNumber | 判断SD大模型版本是1.5还是xl | +| ListWrapper | 包装成列表（任意类型） | +| BboxToCropData | bbox转cropData，方便接入was节点使用 | +| BboxToBbox | bbox两种格式(x,y,w,h)和(x1,y1,x2,y2)的相互转换 | +| BboxesToBboxes | BboxToBbox节点的列表版本 | +| SelectBbox | 从Bbox列表中选择一个 | +| SelectBboxes | 从Bbox列表中选择多个 | +| CropImageByBbox | 根据Bbox区域裁剪图片 | +| MaskByBboxes | 根据Bbox列表画遮罩 | Tips: base64格式字符串比较长，会导致界面卡顿，接口请求带宽可能也会有瓶颈，条件允许可以把图片上传到OSS服务器得到URL，然后用LoadImageFromUrl加载，由于无相关OSS账号，上传OSS节点需自行编写，暂不支持。 ### [示例](example/example.png) ![save api extended](docs/example_note.png) + ![save api extended](example/example_1.png) ## 更新记录 +### 2024-09-04 +- 添加一些bbox相关节点 + ### 2024-08-08 - 菜单适配ComfyUI前端新界面 diff --git a/easyapi/BboxNode.py b/easyapi/BboxNode.py new file mode 100644 index 0000000..01403e7 --- /dev/null +++ b/easyapi/BboxNode.py @@ -0,0 +1,317 @@ +from json import JSONDecoder + +import torch + +from .util import any_type + + +class BboxToCropData: + @classmethod + def INPUT_TYPES(self): + return {"required": { + "bbox": ("BBOX", {"forceInput": True}), + }, + } + + RETURN_TYPES = (any_type,) + RETURN_NAMES = ("crop_data", ) + + FUNCTION = "convert" + + OUTPUT_NODE = False + CATEGORY = "EasyApi/Bbox" + + INPUT_IS_LIST = False + OUTPUT_IS_LIST = (False, ) + DESCRIPTION = "可以把bbox(x,y,w,h)转换为crop_data((w,h),(x,y,x+w,y+w)，配合was节点使用" + + def convert(self, bbox): + x, y, w, h = bbox + return (((w, h), (x, y, x+w, y+h),),) + + +class BboxToCropData: + @classmethod + def INPUT_TYPES(self): + return {"required": { + "bbox": ("BBOX", {"forceInput": True}), + }, + "optional": { + "is_xywh": ("BOOLEAN", {"default": False}), + } + } + + RETURN_TYPES = (any_type,) + RETURN_NAMES = ("crop_data", ) + + FUNCTION = "convert" + + OUTPUT_NODE = False + CATEGORY = "EasyApi/Bbox" + + INPUT_IS_LIST = False + OUTPUT_IS_LIST = (False, ) + DESCRIPTION = "可以把bbox(x,y,w,h)转换为crop_data((w,h),(x,y,x+w,y+w)，配合was节点使用\nis_xywh表示bbox的格式是(x,y,w,h)还是(x,y,x1,y1)。" + + def convert(self, bbox, is_xywh=False): + if is_xywh: + x, y, w, h = bbox + else: + x, y, x_1, y_1 = bbox + w = x_1 - x + h = y_1 - y + return (((w, h), (x, y, x+w, y+h),),) + + +class BboxToBbox: + @classmethod + def INPUT_TYPES(self): + return {"required": { + "bbox": ("BBOX", {"forceInput": True}), + }, + "optional": { + "is_xywh": ("BOOLEAN", {"default": False}), + "to_xywh": ("BOOLEAN", {"default": False}) + } + } + + RETURN_TYPES = (any_type,) + RETURN_NAMES = ("bbox", ) + + FUNCTION = "convert" + + OUTPUT_NODE = False + CATEGORY = "EasyApi/Bbox" + + INPUT_IS_LIST = False + OUTPUT_IS_LIST = (False, ) + DESCRIPTION = "可以把bbox转换为(x1,y1,x2,y2)或(x,y,w,h)，返回任意类型，配合其它bbox节点使用\n is_xywh表示输入的bbox的格式是(x,y,w,h)还是(x,y,x1,y1)。\n to_xywh表示返回的bbox的格式是(x,y,w,h)还是(x,y,x1,y1)。" + + def convert(self, bbox, is_xywh=False, to_xywh=False): + if is_xywh: + x, y, w, h = bbox + else: + x, y, x_1, y_1 = bbox + w = x_1 - x + h = y_1 - y + if to_xywh: + return ((x, y, w, h),) + else: + return ((x, y, x+w, y+h),) + + +class BboxesToBboxes: + @classmethod + def INPUT_TYPES(self): + return {"required": { + "bboxes": ("BBOX", {"forceInput": True}), + }, + "optional": { + "is_xywh": ("BOOLEAN", {"default": False}), + "to_xywh": ("BOOLEAN", {"default": False}) + } + } + + RETURN_TYPES = (any_type,) + RETURN_NAMES = ("bbox", ) + + FUNCTION = "convert" + + OUTPUT_NODE = False + CATEGORY = "EasyApi/Bbox" + + INPUT_IS_LIST = False + OUTPUT_IS_LIST = (False, ) + DESCRIPTION = "可以把bbox转换为(x1,y1,x2,y2)或(x,y,w,h)，返回任意类型，配合其它bbox节点使用\n is_xywh表示输入的bbox的格式是(x,y,w,h)还是(x,y,x1,y1)。\n to_xywh表示返回的bbox的格式是(x,y,w,h)还是(x,y,x1,y1)。" + + def convert(self, bboxes, is_xywh=False, to_xywh=False): + new_bboxes = list() + for bbox in bboxes: + if is_xywh: + x, y, w, h = bbox + else: + x, y, x_1, y_1 = bbox + w = x_1 - x + h = y_1 - y + if to_xywh: + new_bboxes.append((x, y, w, h)) + else: + new_bboxes.append((x, y, x+w, y+h)) + return (new_bboxes,) + + +class SelectBbox: + def __init__(self): + self.models = {} + + @classmethod + def INPUT_TYPES(self): + return { + "required": { + "index": ('INT', {'default': 0, 'step': 1, 'min': 0, 'max': 50}), + }, + "optional": { + "bboxes": ('BBOX', {'forceInput': True}), + "bboxes_json": ('STRING', {'forceInput': True}), + } + } + + RETURN_TYPES = ("BBOX",) + RETURN_NAMES = ("bbox",) + + FUNCTION = "select" + + OUTPUT_NODE = False + CATEGORY = "EasyApi/Bbox" + + # INPUT_IS_LIST = False + # OUTPUT_IS_LIST = (False, False) + + def select(self, index, bboxes=None, bboxes_json=None): + if bboxes is None: + if bboxes_json is not None: + _bboxes = JSONDecoder().decode(bboxes_json) + if len(_bboxes) > index: + return (_bboxes[index], ) + if isinstance(bboxes, list) and len(bboxes) > index: + return (bboxes[index], ) + return (None, ) + + +class SelectBboxes: + def __init__(self): + self.models = {} + + @classmethod + def INPUT_TYPES(self): + return { + "required": { + "index": ('STRING', {'default': "0"}), + }, + "optional": { + "bboxes": ('BBOX', {'forceInput': True}), + "bboxes_json": ('STRING', {'forceInput': True}), + } + } + + RETURN_TYPES = ("BBOX",) + RETURN_NAMES = ("bboxes",) + + FUNCTION = "select" + + OUTPUT_NODE = False + CATEGORY = "EasyApi/Bbox" + + # INPUT_IS_LIST = False + # OUTPUT_IS_LIST = (False, False) + DESCRIPTION = "根据索引(多个逗号分隔)选择bbox" + + def select(self, index, bboxes=None, bboxes_json=None): + indices = [int(i) for i in index.split(",")] + if bboxes is None: + if bboxes_json is not None: + _bboxes = JSONDecoder().decode(bboxes_json) + filtered_bboxes = [_bboxes[i] for i in indices if 0 <= i < len(_bboxes)] + return (filtered_bboxes, ) + if isinstance(bboxes, list): + filtered_bboxes = [bboxes[i] for i in indices if 0 <= i < len(bboxes)] + return (filtered_bboxes,) + return (None, ) + + +class CropImageByBbox: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "bbox": ("BBOX",), + "margin": ("INT", {"default": 16}), + } + } + + RETURN_TYPES = ("IMAGE", "MASK", "BBOX") + RETURN_NAMES = ("crop_image", "mask", "crop_bbox") + FUNCTION = "crop" + CATEGORY = "EasyApi/Bbox" + + def crop(self, image: torch.Tensor, bbox, margin): + x, y, x1, y1 = bbox + w = x1 - x + h = y1 - y + image_height = image.shape[1] + image_width = image.shape[2] + # 左上角坐标 + x = min(x, image_width) + y = min(y, image_height) + # 右下角坐标 + to_x = min(w + x + margin, image_width) + to_y = min(h + y + margin, image_height) + # 防止越界 + x = max(0, x - margin) + y = max(0, y - margin) + to_x = max(0, to_x) + to_y = max(0, to_y) + # 按区域截取图片 + crop_img = image[:, y:to_y, x:to_x, :] + new_bbox = (x, y, to_x, to_y) + # 创建与image相同大小的全零张量作为遮罩 + mask = torch.zeros((image_height, image_width), dtype=torch.uint8) # 使用uint8类型 + # 在mask上设置new_bbox区域为1 + mask[new_bbox[1]:new_bbox[3], new_bbox[0]:new_bbox[2]] = 1 + # 如果需要转换为浮点数，并且增加一个通道维度, 形状变为 (1, height, width) + mask_tensor = mask.unsqueeze(0) + return crop_img, mask_tensor, new_bbox, + + +class MaskByBboxes: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "bboxes": ("BBOX",), + } + } + + RETURN_TYPES = ("MASK", ) + RETURN_NAMES = ("mask", ) + FUNCTION = "crop" + CATEGORY = "EasyApi/Bbox" + DESCRIPTION = "根据bboxes生成遮罩, bboxes格式是(x, y, w, h)" + + def crop(self, image: torch.Tensor, bboxes): + image_height = image.shape[1] + image_width = image.shape[2] + + # 创建与image相同大小的全零张量作为遮罩 + mask = torch.zeros((image_height, image_width), dtype=torch.uint8) + # 在mask上设置new_bbox区域为1 + for bbox in bboxes: + x, y, w, h = bbox + mask[y:y+h, x:x+w] = 1 + # 如果需要转换为浮点数，并且增加一个通道维度, 形状变为 (1, height, width) + mask_tensor = mask.unsqueeze(0) + return mask_tensor, + + +NODE_CLASS_MAPPINGS = { + "BboxToCropData": BboxToCropData, + "BboxToBbox": BboxToBbox, + "BboxesToBboxes": BboxesToBboxes, + "SelectBbox": SelectBbox, + "SelectBboxes": SelectBboxes, + "CropImageByBbox": CropImageByBbox, + "MaskByBboxes": MaskByBboxes, +} + + +NODE_DISPLAY_NAME_MAPPINGS = { + "BboxToCropData": "BboxToCropData", + "BboxToBbox": "BboxToBbox", + "BboxesToBboxes": "BboxesToBboxes", + "SelectBbox": "SelectBbox", + "SelectBboxes": "SelectBboxes", + "CropImageByBbox": "CropImageByBbox", + "MaskByBboxes": "MaskByBboxes", +} diff --git a/easyapi/DetectNode.py b/easyapi/DetectNode.py index a779218..1ab28ee 100644 --- a/easyapi/DetectNode.py +++ b/easyapi/DetectNode.py @@ -1,7 +1,7 @@ import os from PIL import Image -from json import JSONEncoder, JSONDecoder +from json import JSONEncoder import numpy as np from .util import tensor_to_pil, pil_to_tensor, hex_to_rgba diff --git a/easyapi/UtilNode.py b/easyapi/UtilNode.py index 3c369a9..8144a92 100644 --- a/easyapi/UtilNode.py +++ b/easyapi/UtilNode.py @@ -2,7 +2,7 @@ import torch from comfy.model_patcher import ModelPatcher import comfy.model_base -from .util import tensor_to_pil, hex_to_rgba +from .util import tensor_to_pil, hex_to_rgba, any_type class GetImageBatchSize: @@ -362,6 +362,39 @@ class SDBaseVerNumber: return (*get_model_sd_version(model),) +class ListWrapper: + @classmethod + def INPUT_TYPES(self): + return {"required": { + "any_1": (any_type, {"forceInput": True}), + }, + "optional": { + "any_2": (any_type, {"forceInput": True}), + }, + } + + RETURN_TYPES = (any_type,) + # RETURN_NAMES = ("STRING", ) + + FUNCTION = "to_list" + + OUTPUT_NODE = False + CATEGORY = "EasyApi/List" + + INPUT_IS_LIST = False + OUTPUT_IS_LIST = (False, ) + DESCRIPTION = "把输入放到一个列表中，如bbox转bboxes" + + def to_list(self, any_1, any_2=None): + if any_1 is None: + return None + else: + if any_2 is None: + return ((any_1,),) + else: + return ((any_1, any_2),) + + NODE_CLASS_MAPPINGS = { "GetImageBatchSize": GetImageBatchSize, "JoinList": JoinList, @@ -377,6 +410,7 @@ NODE_CLASS_MAPPINGS = { "ColorPicker": ColorPicker, "ImageEqual": ImageEqual, "SDBaseVerNumber": SDBaseVerNumber, + "ListWrapper": ListWrapper, } # A dictionary that contains the friendly/humanly readable titles for the nodes @@ -395,4 +429,5 @@ NODE_DISPLAY_NAME_MAPPINGS = { "ColorPicker": "Color Picker", "ImageEqual": "Image Equal Judgment", "SDBaseVerNumber": "SD Base Version Number", + "ListWrapper": "ListWrapper", } diff --git a/easyapi/util.py b/easyapi/util.py index 414672e..c9eba68 100644 --- a/easyapi/util.py +++ b/easyapi/util.py @@ -69,3 +69,13 @@ def hex_to_rgba(hex_color): else: a = 255 return r, g, b, a + + +class AnyType(str): + """A special class that is always equal in not equal comparisons. Credit to pythongosssss""" + + def __ne__(self, __value: object) -> bool: + return False + + +any_type = AnyType("*") diff --git a/example/example_1.png b/example/example_1.png new file mode 100644 index 0000000..d7bd3d6 Binary files /dev/null and b/example/example_1.png differ diff --git a/pyproject.toml b/pyproject.toml index 348d007..998f613 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui-easyapi-nodes" description = "Provides some features and nodes related to API calls." -version = "1.0.1" +version = "1.0.2" license = { file = "LICENSE" } dependencies = ["segment_anything", "simple_lama_inpainting", "insightface"]