xinyun/ComfyUI-KJNodes
1
0
Fork 0
mirror of https://git.datalinker.icu/kijai/ComfyUI-KJNodes.git synced 2025-12-13 06:44:29 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add bbox_format option for bbox visualize node

Browse Source
This commit is contained in:
kijai 2025-08-15 12:46:02 +03:00
parent 87d0cf42db
commit 2f7300dc54
1 changed files with 18 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
nodes/batchcrop_nodes.py
View File

@ -694,6 +694,7 @@ class BboxVisualize:
"images": ("IMAGE",), "images": ("IMAGE",),
"bboxes": ("BBOX",), "bboxes": ("BBOX",),
"line_width": ("INT", {"default": 1,"min": 1, "max": 10, "step": 1}), "line_width": ("INT", {"default": 1,"min": 1, "max": 10, "step": 1}),
"bbox_format": (["xywh", "xyxy"], {"default": "xywh"}),
}, },
} }
@ -706,26 +707,33 @@ Visualizes the specified bbox on the image.
CATEGORY = "KJNodes/masking" CATEGORY = "KJNodes/masking"
def visualizebbox(self, bboxes, images, line_width): def visualizebbox(self, bboxes, images, line_width, bbox_format):
image_list = [] image_list = []
for image, bbox in zip(images, bboxes): for image, bbox in zip(images, bboxes):
x_min, y_min, width, height = bbox if bbox_format == "xywh":
x_min, y_min, width, height = bbox
elif bbox_format == "xyxy":
x_min, y_min, x_max, y_max = bbox
width = x_max - x_min
height = y_max - y_min
else:
raise ValueError(f"Unknown bbox_format: {bbox_format}")
# Ensure bbox coordinates are integers # Ensure bbox coordinates are integers
x_min = int(x_min) x_min = int(x_min)
y_min = int(y_min) y_min = int(y_min)
width = int(width) width = int(width)
height = int(height) height = int(height)
# Permute the image dimensions # Permute the image dimensions
image = image.permute(2, 0, 1) image = image.permute(2, 0, 1)
# Clone the image to draw bounding boxes # Clone the image to draw bounding boxes
img_with_bbox = image.clone() img_with_bbox = image.clone()
# Define the color for the bbox, e.g., red # Define the color for the bbox, e.g., red
color = torch.tensor([1, 0, 0], dtype=torch.float32) color = torch.tensor([1, 0, 0], dtype=torch.float32)
# Ensure color tensor matches the image channels # Ensure color tensor matches the image channels
if color.shape[0] != img_with_bbox.shape[0]: if color.shape[0] != img_with_bbox.shape[0]:
color = color.unsqueeze(1).expand(-1, line_width) color = color.unsqueeze(1).expand(-1, line_width)
@ -735,23 +743,21 @@ Visualizes the specified bbox on the image.
# Top horizontal line # Top horizontal line
if y_min + lw < img_with_bbox.shape[1]: if y_min + lw < img_with_bbox.shape[1]:
img_with_bbox[:, y_min + lw, x_min:x_min + width] = color[:, None] img_with_bbox[:, y_min + lw, x_min:x_min + width] = color[:, None]
# Bottom horizontal line # Bottom horizontal line
if y_min + height - lw < img_with_bbox.shape[1]: if y_min + height - lw < img_with_bbox.shape[1]:
img_with_bbox[:, y_min + height - lw, x_min:x_min + width] = color[:, None] img_with_bbox[:, y_min + height - lw, x_min:x_min + width] = color[:, None]
# Left vertical line # Left vertical line
if x_min + lw < img_with_bbox.shape[2]: if x_min + lw < img_with_bbox.shape[2]:
img_with_bbox[:, y_min:y_min + height, x_min + lw] = color[:, None] img_with_bbox[:, y_min:y_min + height, x_min + lw] = color[:, None]
# Right vertical line # Right vertical line
if x_min + width - lw < img_with_bbox.shape[2]: if x_min + width - lw < img_with_bbox.shape[2]:
img_with_bbox[:, y_min:y_min + height, x_min + width - lw] = color[:, None] img_with_bbox[:, y_min:y_min + height, x_min + width - lw] = color[:, None]
# Permute the image dimensions back # Permute the image dimensions back
img_with_bbox = img_with_bbox.permute(1, 2, 0).unsqueeze(0) img_with_bbox = img_with_bbox.permute(1, 2, 0).unsqueeze(0)
image_list.append(img_with_bbox) image_list.append(img_with_bbox)
return (torch.cat(image_list, dim=0),)
return (torch.cat(image_list, dim=0),) return (torch.cat(image_list, dim=0),)