Skip to content

Commit

Permalink
add some nodes aboute bbox
Browse files Browse the repository at this point in the history
  • Loading branch information
刘雪峰 committed Sep 4, 2024
1 parent c11ff77 commit 966ddc4
Show file tree
Hide file tree
Showing 7 changed files with 377 additions and 3 deletions.
12 changes: 12 additions & 0 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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前端新界面

Expand Down
317 changes: 317 additions & 0 deletions easyapi/BboxNode.py
View file
Original file line number Diff line number Diff line change
@@ -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",
}
2 changes: 1 addition & 1 deletion easyapi/DetectNode.py
View file
Original file line number Diff line number Diff line change
@@ -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
Expand Down
Loading

0 comments on commit 966ddc4

Please sign in to comment.