318 lines
9.0 KiB
Python
318 lines
9.0 KiB
Python
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",
|
|
}
|