409 lines
13 KiB
Python
409 lines
13 KiB
Python
from json import JSONDecoder
|
|
|
|
import torch
|
|
|
|
import nodes
|
|
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:
|
|
@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)
|
|
DESCRIPTION = "根据索引过滤"
|
|
|
|
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:
|
|
@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 = "根据索引(支持逗号分隔)过滤"
|
|
|
|
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, "tooltip": "bbox矩形区域向外扩张的像素距离"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK", "BBOX", "INT", "INT")
|
|
RETURN_NAMES = ("crop_image", "mask", "crop_bbox", "w", "h")
|
|
FUNCTION = "crop"
|
|
CATEGORY = "EasyApi/Bbox"
|
|
DESCRIPTION = "根据bbox区域裁剪图片。 bbox的格式是左上角和右下角坐标： [x,y,x1,y1]"
|
|
|
|
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, to_x - x, to_y - y,
|
|
|
|
|
|
class CropTargetSizeImageByBbox:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"bbox": ("BBOX",{"forceInput": True, "tooltip": "参考区域坐标"}),
|
|
"width": ("INT", {"default": 512, "min": 1, "max": nodes.MAX_RESOLUTION, "step": 1, "tooltip": "目标宽度"}),
|
|
"height": ("INT", {"default": 512, "min": 1, "max": nodes.MAX_RESOLUTION, "step": 1, "tooltip": "目标高度"}),
|
|
"contain": ("BOOLEAN", {"default": False, "tooltip": "是否始终包含bbox完整区域"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK", "BBOX", "INT", "INT")
|
|
RETURN_NAMES = ("crop_image", "mask", "crop_bbox", "w", "h")
|
|
FUNCTION = "crop"
|
|
CATEGORY = "EasyApi/Bbox"
|
|
DESCRIPTION = "根据bbox区域中心裁剪指定大小图片。 bbox的格式是左上角和右下角坐标： [x,y,x1,y1]"
|
|
|
|
def calc_area(self, image_width, image_height, rect_top_left, rect_bottom_right, w, h):
|
|
"""
|
|
以给定的矩形中心点为中心计算指定宽高的矩形框坐标
|
|
Args:
|
|
image_width: 图片高度
|
|
image_height: 图片宽度
|
|
rect_top_left: 矩形框左上角坐标
|
|
rect_bottom_right: 矩形框右下角坐标
|
|
w: 目标宽度
|
|
h: 目标高度
|
|
|
|
Returns:
|
|
|
|
"""
|
|
# 计算矩形的宽和高
|
|
x, y = rect_top_left
|
|
x1, y1 = rect_bottom_right
|
|
|
|
# 否则，计算矩形的中心(取整)
|
|
center_x = (x + x1) // 2
|
|
center_y = (y + y1) // 2
|
|
left_w = w // 2
|
|
right_w = w - left_w
|
|
top_h = h // 2
|
|
bottom_h = h - top_h
|
|
|
|
# 计算新的坐标
|
|
new_top_left_x = max(0, center_x - left_w)
|
|
new_top_left_y = max(0, center_y - top_h)
|
|
new_bottom_right_x = min(image_width, center_x + right_w)
|
|
new_bottom_right_y = min(image_height, center_y + bottom_h)
|
|
|
|
# 如果坐标越界，调整坐标
|
|
if new_top_left_x == 0:
|
|
# 左边可能超过边界了，尝试把左边超出部分加到右边
|
|
new_bottom_right_x = min(image_width, new_bottom_right_x + (left_w - center_x))
|
|
elif new_bottom_right_x == image_width:
|
|
# 右边可能超过边界了，尝试把右边超出部分加到左边
|
|
new_top_left_x = max(0, new_top_left_x - (center_x + left_w - image_width))
|
|
|
|
if new_top_left_y == 0:
|
|
# 上边可能超过边界了，尝试把上边超出部分加到下边
|
|
new_bottom_right_y = min(image_height, new_bottom_right_y + (top_h - center_y))
|
|
elif new_bottom_right_y == image_height:
|
|
# 下边可能超过边界了，尝试把下边超出部分加到上边
|
|
new_top_left_y = max(0, new_top_left_y - (center_y + top_h - image_height))
|
|
|
|
return new_top_left_x, new_top_left_y, new_bottom_right_x, new_bottom_right_y
|
|
|
|
def crop(self, image: torch.Tensor, bbox, width, height, contain):
|
|
x, y, x1, y1 = bbox
|
|
image_height = image.shape[1]
|
|
image_width = image.shape[2]
|
|
|
|
new_x, new_y, to_x, to_y = self.calc_area(image_width, image_height, (x, y), (x1, y1), width, height)
|
|
|
|
if contain:
|
|
new_x = min(new_x, x)
|
|
new_y = min(new_y, y)
|
|
to_x = max(to_x, x1)
|
|
to_y = max(to_y, y1)
|
|
# 按区域截取图片
|
|
crop_img = image[:, new_y:to_y, new_x:to_x, :]
|
|
new_bbox = (new_x, new_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, to_x - new_x, to_y - new_y,
|
|
|
|
|
|
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,
|
|
"CropTargetSizeImageByBbox": CropTargetSizeImageByBbox,
|
|
}
|
|
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"BboxToCropData": "BboxToCropData",
|
|
"BboxToBbox": "BboxToBbox",
|
|
"BboxesToBboxes": "BboxesToBboxes",
|
|
"SelectBbox": "SelectBbox",
|
|
"SelectBboxes": "SelectBboxes",
|
|
"CropImageByBbox": "CropImageByBbox",
|
|
"MaskByBboxes": "MaskByBboxes",
|
|
"CropTargetSizeImageByBbox": "CropTargetSizeImageByBbox",
|
|
}
|