From 6c9f883ee455e22eee24b4eb76e7614186dc31c2 Mon Sep 17 00:00:00 2001 From: Yury Krychala Date: Wed, 8 Apr 2026 04:50:06 -0400 Subject: [PATCH] Added schedulers, updated to V3 schema, refactored code --- README.md | 11 ++ __init__.py | 25 ++-- nodes.py | 152 ------------------------ nodes/image_nodes.py | 255 ++++++++++++++++++++++++++++------------ nodes/schedule_nodes.py | 158 +++++++++++++++++++++++++ nodes/utils.py | 77 ++++++++++++ requirements.txt | 5 + 7 files changed, 437 insertions(+), 246 deletions(-) delete mode 100644 nodes.py create mode 100644 nodes/schedule_nodes.py create mode 100644 nodes/utils.py diff --git a/README.md b/README.md index bdbfeb8..8d6c12f 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,13 @@ # ComfyUI-BBoxNodes BBoxes (BoundingBoxes) handling and visualization nodes + +Documentation is mostly in the node descriptions and tooltips. + +# Installation +1. Clone this repo into `custom_nodes` folder. +2. Install dependencies: `pip install -r requirements.txt` + or if you use the portable install, run this in ComfyUI_windows_portable -folder: + + `python_embeded\python.exe -m pip install -r ComfyUI\custom_nodes\ComfyUI-BBoxNodes\requirements.txt` + + diff --git a/__init__.py b/__init__.py index a03f0b8..8994498 100644 --- a/__init__.py +++ b/__init__.py @@ -1,22 +1,11 @@ from .nodes.image_nodes import * +from .nodes.schedule_nodes import * +from comfy_api.latest import ComfyExtension, io -NODE_CONFIG = { - "BboxesVisualize": {"class": BboxesVisualize, "name": "Bounding Boxes Visualize"}, -} +class BboxNodesExtension(ComfyExtension): + async def get_node_list(self) -> list[type[io.ComfyNode]]: + return [BboxesVisualize, TrackingVisualize, TrackingToBBoxScheduler, BBoxScheduler] -def generate_node_mappings(node_config): - node_class_mappings = {} - node_display_name_mappings = {} +async def comfy_entrypoint() -> ComfyExtension: + return BboxNodesExtension() - for node_name, node_info in node_config.items(): - node_class_mappings[node_name] = node_info["class"] - node_display_name_mappings[node_name] = node_info.get("name", node_info["class"].__name__) - - return node_class_mappings, node_display_name_mappings - - -# Map all your custom nodes classes with the names that will be displayed in the UI. -NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS = generate_node_mappings(NODE_CONFIG) - - -__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/nodes.py b/nodes.py deleted file mode 100644 index 9f507e2..0000000 --- a/nodes.py +++ /dev/null @@ -1,152 +0,0 @@ -# Package Modules -import os -from typing import Union, BinaryIO, Dict, List, Tuple, Optional -import time - -# ComfyUI Modules -import folder_paths -from comfy.utils import ProgressBar - -# Your Modules -from .modules.calculator import CalculatorModel - - -# Basic practice to get paths from ComfyUI -custom_nodes_script_dir = os.path.dirname(os.path.abspath(__file__)) -custom_nodes_model_dir = os.path.join(folder_paths.models_dir, "my-custom-nodes") -custom_nodes_output_dir = os.path.join(folder_paths.get_output_directory(), "my-custom-nodes") - - -# These are example nodes that only contains basic functionalities with some comments. -# If you need detailed explanation, please refer to : https://docs.comfy.org/essentials/custom_node_walkthrough -# First Node: -class MyModelLoader: - # Define the input parameters of the node here. - @classmethod - def INPUT_TYPES(s): - my_models = ["Model A", "Model B", "Model C"] - - return { - # If the key is "required", the value must be filled. - "required": { - # `my_models` is the list, so it will be shown as a dropdown menu in the node. ( So that user can select one of them. ) - # You must provide the value in the tuple format. e.g. ("value",) or (3,) or ([1, 2],) etc. - "model": (my_models,), - "device": (['cuda', 'cpu', 'auto'],), - }, - # If the key is "optional", the value is optional. - "optional": { - "compute_type": (['float32', 'float16'],), - } - } - - # Define these constants inside the node. - # `RETURN_TYPES` is important, as it limits the parameter types that can be passed to the next node, in `INPUT_TYPES()` above. - RETURN_TYPES = ("MY_MODEL",) - RETURN_NAMES = ("my_model",) - # `FUNCTION` is the function name that will be called in the node. - FUNCTION = "load_model" - # `CATEGORY` is the category name that will be used when user searches the node. - CATEGORY = "CustomNodesTemplate" - - # In the function, use same parameter names as you specified in `INPUT_TYPES()` - def load_model(self, - model: str, - device: str, - compute_type: Optional[str] = None, - ) -> Tuple[CalculatorModel]: - calculator_model = CalculatorModel() - calculator_model.load_model(model, device, compute_type) - - # You can use `comfy.utils.ProgressBar` to show the progress of the process. - # First, initialize the total amount of the process. - total_steps = 5 - comfy_pbar = ProgressBar(total_steps) - # Then, update the progress. - for i in range(1, total_steps): - time.sleep(1) - comfy_pbar.update(i) # Alternatively, you can use `comfy_pbar.update_absolute(value)` to update the progress with absolute value. - - # Return the model as a tuple. - return (calculator_model, ) - - -# Second Node -class CalculatePlus: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "model": ("MY_MODEL", ), - }, - # Specify the parameters with type and default value. - "optional": { - "a": ("INT", {"default": 5}), - "b": ("INT", {"default": 10}), - } - } - - RETURN_TYPES = ("INT",) - RETURN_NAMES = ("plus_value",) - FUNCTION = "plus" - CATEGORY = "CustomNodesTemplate" - - def plus(self, - model: CalculatorModel, - a: Optional[int], - b: Optional[int], - ) -> Tuple[int]: - result = model.plus(a, b) - return (result, ) - - - -# Third Node -class CalculateMinus: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "model": ("MY_MODEL", ), - "a": ("INT", ), - }, - "optional": { - "b": ("INT", {"default": 10}), - } - } - - RETURN_TYPES = ("INT",) - RETURN_NAMES = ("minus_value",) - FUNCTION = "minus" - CATEGORY = "CustomNodesTemplate" - - def minus(self, - model: CalculatorModel, - a: Optional[int], - b: Optional[int], - ) -> Tuple[int]: - result = model.minus(a, b) - return (result, ) - - - -# Output Node -class ExampleOutputNode: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "value": ("INT", ), - }, - } - - # If the node is output node, set this to True. - OUTPUT_NODE = True - RETURN_TYPES = ("INT",) - RETURN_NAMES = ("int",) - FUNCTION = "result" - CATEGORY = "CustomNodesTemplate" - - def result(self, - value: int,) -> Tuple[int]: - return (value, ) diff --git a/nodes/image_nodes.py b/nodes/image_nodes.py index 020d595..0d1d0a8 100644 --- a/nodes/image_nodes.py +++ b/nodes/image_nodes.py @@ -1,109 +1,212 @@ -from comfy_api.latest import io import torch +from utils import push_bbox, parse_tracking, body_part_names, draw_frame_index +from torchvision import transforms +from comfy.utils import ProgressBar +from comfy_api.latest import io +from PIL import ImageDraw -def push_bbox(x, y, width, height, draw_boxes): - bbox = [x, y, width, height] - draw_boxes.append(bbox) +colors = [ + [1, 0, 0], + [0, 1, 0], + [0, 0, 1], + [0.5, 0.5, 0], + [0, 0.5, 0.5], + [0.5, 0, 0.5], + [0.25, 0.75, 0], + [0.75, 0.25, 0], + [0, 0.25, 0.75], + [0, 0.75, 0.25] +] -class BboxesVisualize: +class BboxesVisualize(io.ComfyNode): - def define_schema(self): + + @classmethod + def define_schema(cls) -> io.Schema: return io.Schema( - node_id="VisualizeBBoxes", + node_id="BboxesVisualize", + display_name="Bounding Boxes Visualize", category="BBoxNodes", - search_aliases=["Visualize", "bbox", "bounding box"], - description="Visualizes the provided bboxes on the image for specified person. Specify 0 for all persons", + description = """ + Visualizes the provided bboxes on the image for specified person(s). Specify 0 for all persons. + Not compatible with KJNodes bboxes + """, inputs=[ - io.Image.Input("image"), - io.BoundingBox.Input("bboxes", force_input=True), - io.Int.Input("line_width", default=0, min=0, max=10, step=1, - tooltip="Width of the line"), - io.Int.Input("person_index", default=0, min=0, max=5, step=1, - tooltip="Index of the person. 0 is all persons"), + io.Image.Input("images", + optional=False, + tooltip="The input images to process" + ), + io.BoundingBox.Input("bboxes", + force_input=True, + tooltip="Bounding boxes (not compatible with KJNodes)" + ), + io.Int.Input("line_width", + min=1, + max=5, + step=1, + default=1, + optional=False, + tooltip="Rectangle border width" + ), + io.Int.Input("person_index", + min=0, + max=10, + step=1, + default=0, + optional=False, + tooltip="The person index on the image starting from 1. 0 for all persons" + ), ], outputs=[ - io.Image.Output(tooltip="Image with boxes drawn on the image"), + io.Image.Output("images", + tooltip="Output images with drawn bounding boxes"), ], ) @classmethod - def execute(cls, image, bboxes, line_width, person_index) -> io.NodeOutput: + def execute(cls, images, bboxes, line_width, person_index) -> io.NodeOutput: - draw_boxes = [] image_list = [] - total_frames = image.shape[0] - img_h = image.shape[1] - img_w = image.shape[2] + if images.shape[0] == 0: + raise ValueError("Input image batch is empty") if not isinstance(bboxes, list): bboxes = [[bboxes]] elif len(bboxes) == 0: - return io.NodeOutput(image) + return io.NodeOutput(images) - for frame_idx in range(total_frames): - frame_bboxes = bboxes[min(frame_idx, len(bboxes) - 1)] - if not frame_bboxes: - continue + steps = images.shape[0] + pbar = ProgressBar(steps) + batch_count = 0 - if person_index == 0: - for b in frame_bboxes: - x = b["x"] - y = b["y"] - width = b["width"] - height = b["height"] + for image, frame_bboxes in zip(images, bboxes): + + draw_boxes = [] + for index, box in enumerate(frame_bboxes): + if person_index == 0 or index == person_index - 1: + x = box["x"] + y = box["y"] + width = box["width"] + height = box["height"] push_bbox(x, y, width, height, draw_boxes) - else: - index = 0 - if person_index > 0 and len(frame_bboxes) <= person_index: - index = person_index - 1 - bbox = frame_bboxes[index] - x = bbox["x"] - y = bbox["y"] - width = bbox["width"] - height = bbox["height"] - push_bbox(x, y, width, height, draw_boxes) - - for bbox in draw_boxes: - - x1 = min(bbox[0], img_w - 1) - y1 = min(bbox[1], img_h - 1) - x2 = min(img_w, bbox[0] + bbox[2]) - y2 = min(img_h, bbox[1] + bbox[3]) # Permute the image dimensions - image = image.permute(2, 0, 1) + img_with_bbox = image.permute(2, 0, 1) + pil_image = transforms.ToPILImage()(img_with_bbox) + draw = ImageDraw.Draw(pil_image) - # Clone the image to draw bounding boxes - img_with_bbox = image.clone() + # go through selected draw boxes + for color_index, bbox in enumerate(draw_boxes): - # Define the color for the bbox, e.g., red - color = torch.tensor([1, 0, 0], dtype=torch.float32) + # Define the color for the bbox, e.g., red + color = tuple(int(255 * x) for x in colors[color_index])[:3] - # Ensure color tensor matches the image channels - if color.shape[0] != img_with_bbox.shape[0]: - color = color.unsqueeze(1).expand(-1, line_width) + # Draw the bbox rectangle + x, y, width, height = bbox + draw.rectangle([x, y, x + width, y + width], outline=color, width=line_width) - # Draw lines for each side of the bbox with the specified line width - for lw in range(line_width): - # Top horizontal line - if y1 + lw < img_h: - img_with_bbox[:, y1 + lw, x1:x2] = color[:, None] - - # Bottom horizontal line - if y2 - lw < img_h: - img_with_bbox[:, y2 - lw, x1:x2] = color[:, None] - - # Left vertical line - if x1 + lw < img_w: - img_with_bbox[:, y1:y2, x1 + lw] = color[:, None] - - # Right vertical line - if x2 - lw < img_w: - img_with_bbox[:, y1:y2, x2 - lw] = color[:, None] + #Drawing frame number in the top left corner of the frame + draw_frame_index(batch_count, draw) # Permute the image dimensions back - img_with_bbox = img_with_bbox.permute(1, 2, 0).unsqueeze(0) + img_with_bbox = transforms.ToTensor()(pil_image).permute(1, 2, 0) image_list.append(img_with_bbox) + pbar.update(batch_count) + batch_count += 1 - return io.NodeOutput(image_list) \ No newline at end of file + return io.NodeOutput(torch.stack(image_list).cpu().float()) + +class TrackingVisualize(io.ComfyNode): + + + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="TrackingVisualize", + display_name="Tracking Data Visualize", + category="BBoxNodes", + description = """ + Visualizes the provided tracking info (InstanceDiffusion) on the image for all or specified person. + Similar to KJNode 'DrawInstanceDiffusionTracking'. Compatible with InstanceDiffusion nodes + """, + inputs=[ + io.Image.Input("images", + optional=False, + tooltip="The input images to process" + ), + io.Custom("TRACKING").Input("tracking", + optional=False, + tooltip="Tracking data from InstanceDiffusion" + ), + io.Int.Input("line_width", + min=1, + max=5, + step=1, + default=1, + optional=False, + tooltip="Rectangle border width" + ), + io.Int.Input("person_index", + min=0, + max=10, + step=1, + default=0, + optional=False, + tooltip="The person index on the image starting from 1. 0 for all persons" + ), + io.Combo.Input("body_part", + options=["All", "Head", "Neck", "Shoulder", "LArm", "RArm", "LForearm", "RForearm", "Torso"], + default="All", + optional=False, + tooltip="The body part on the image to draw rectangle for") + ], + outputs=[ + io.Image.Output("images", + tooltip="Output images with drawn bounding boxes"), + ], + ) + + @classmethod + def execute(cls, images, tracking, line_width, person_index, body_part) -> io.NodeOutput: + + image_list = [] + + steps = images.shape[0] + pbar = ProgressBar(steps) + + draw_boxes = {} + # Iterate through body parts in the tracking data + parse_tracking(tracking, body_part, person_index, False, draw_boxes) + + for frame_index,image in enumerate(images): + + # Permute the image dimensions + img_with_bbox = image.permute(2, 0, 1) + pil_image = transforms.ToPILImage()(img_with_bbox) + draw = ImageDraw.Draw(pil_image) + + if draw_boxes.get(frame_index): + bboxes = draw_boxes[frame_index] + + # Draw collected boxes + for bbox in bboxes: + + x1, y1, x2, y2, part = bbox + color_index = body_part_names.index(part) + + # Define the color for the bbox, e.g., red + color = tuple(int(255 * x) for x in colors[color_index])[:3] + + # Draw the tracking box rectangle + draw.rectangle([x1, y1, x2, y2], outline=color, width=line_width) + + #Drawing frame number in the top left corner of the frame + draw_frame_index(frame_index, draw) + + # Permute the image dimensions back + img_with_bbox = transforms.ToTensor()(pil_image).permute(1, 2, 0) + image_list.append(img_with_bbox) + pbar.update(frame_index) + + return io.NodeOutput(torch.stack(image_list).cpu().float()) diff --git a/nodes/schedule_nodes.py b/nodes/schedule_nodes.py new file mode 100644 index 0000000..3c00b33 --- /dev/null +++ b/nodes/schedule_nodes.py @@ -0,0 +1,158 @@ +from comfy_api.latest import io +from utils import parse_tracking, body_part_names, parse_schedule + + +class BBoxScheduler(io.ComfyNode): + + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="BBoxScheduler", + display_name="BBoxes Scheduler", + category="BBoxNodes", + description = """ +Filters out bounding boxes using person index and a scheduler (Optional)' +Not compatible with KJNodes bboxes +""", + inputs=[ + io.BoundingBox.Input("bboxes", + force_input=True, + tooltip="Bounding boxes (not compatible with KJNodes)" + ), + io.Int.Input("person_index", + min=0, + max=10, + step=1, + default=0, + optional=False, + tooltip="The person index on the image starting from 1. 0 for all persons" + ), + io.Boolean.Input("insert_dummy_bbox", + optional=False, + tooltip="Insert dummy one-pixel bbox when all bboxes are filtered out"), + io.String.Input("schedule", + optional=True, + tooltip="Schedule in format 'frame_index:person_index, ...'"), + ], + outputs=[ + io.BoundingBox.Output("bboxes", + tooltip="Filtered out bboxes (Not compatible with KJNodes)"), + ], + ) + + + @classmethod + def execute(cls, bboxes, person_index, insert_dummy_bbox, schedule) -> io.NodeOutput: + if not isinstance(bboxes, list): + bboxes = [[bboxes]] + elif len(bboxes) == 0: + return io.NodeOutput(bboxes) + + result = [[]] * len(bboxes) + frames_schedule = {} + parse_schedule(frames_schedule, schedule) + + person_filter = 0 + for index, bboxes_frame in enumerate(bboxes): + if frames_schedule and frames_schedule.get(str(index)): + person_filter = frames_schedule[str(index)] + + if len(bboxes_frame) == 0: + continue + + result_bbox = [] + + for box_index, box in enumerate(bboxes_frame): + if person_index == 0 or box_index == person_index - 1: + if person_filter == 0 or box_index == person_filter - 1: + result_bbox.append(box) + + if insert_dummy_bbox and result_bbox == []: + result_bbox = [{"x": 0, "y": 0, "width": 1, "height": 1}] + result[index] = result_bbox + + return io.NodeOutput(result) + + + + +class TrackingToBBoxScheduler(io.ComfyNode): + + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="TrackingToBBoxScheduler", + display_name="Convert Tracking To BBoxes", + category="BBoxNodes", + description = """ +Converts the provided tracking info (InstanceDiffusion) to BBoxes format. Can use filtering (Body part, person index) +and a scheduler (Optional). Person index is 0 for all persons, or a specific person starting from 1 +Compatible with InstanceDiffusion nodes. Not compatible with KJNodes bboxes +""", + inputs=[ + io.Custom("TRACKING").Input("tracking", + optional=False, + tooltip="Tracking data from InstanceDiffusion" + ), + io.Int.Input("person_index", + min=0, + max=10, + step=1, + default=0, + optional=False, + tooltip="The person index on the image starting from 1. 0 for all persons" + ), + io.Combo.Input("body_part", + options=["All"] + body_part_names, + default="All", + optional=False, + tooltip="The body part on the image to draw rectangle for"), + io.Boolean.Input("insert_dummy_bbox", + optional=False, + tooltip="Insert dummy one-pixel bbox when all bboxes are filtered out"), + io.String.Input("schedule", + optional=True, + tooltip="Schedule in format 'frame_index:person_index, ...'"), + ], + outputs=[ + io.BoundingBox.Output("bboxes", + tooltip="Output bboxes (Not compatible with KJNodes)"), + ], + ) + + @classmethod + def execute(cls, tracking, person_index, body_part, insert_dummy_bbox, schedule) -> io.NodeOutput: + + frames_schedule = {} + parse_schedule(frames_schedule, schedule) + + tracking_boxes = [] + num_frames = 1 + # Iterate through body parts in the tracking data + parse_tracking(tracking, body_part, person_index, True, tracking_boxes) + + person_filter = 0 + # initialize bboxes + empty_bbox = {"x": 0, "y": 0, "width": 1, "height": 1} + bboxes = [[]] * num_frames + # go through selected tracking boxes + for tracking_box in tracking_boxes: + + index, x1, y1, x2, y2, person = tracking_box + # convert to bbox format + bbox = {"x": x1, "y": y1, "width": x2 - x1, "height": y2 - y1} + if frames_schedule and frames_schedule.get(str(index)): + person_filter = frames_schedule[str(index)] + + if bbox != empty_bbox: + if person_filter == 0 or person_filter == person + 1: + if bboxes[index] != [empty_bbox]: + bboxes[index].append(bbox) + else: + bboxes[index] = [bbox] + + # replace empy bbox with dummy bbox if required + if bboxes[index] == [] and insert_dummy_bbox: + bboxes[index] = [empty_bbox] + + return io.NodeOutput(bboxes) diff --git a/nodes/utils.py b/nodes/utils.py new file mode 100644 index 0000000..b8d75d4 --- /dev/null +++ b/nodes/utils.py @@ -0,0 +1,77 @@ +import folder_paths +import os +from PIL import ImageFont + + +body_part_names = ["Head", "Neck", "Shoulder", "Torso", "RArm", "RForearm", "LArm", "LForearm"] +script_directory = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +folder_paths.add_model_folder_path("bboxnodes_fonts", os.path.join(script_directory, "fonts")) +font_path = folder_paths.get_full_path("bboxnodes_fonts", "FreeMono.ttf") +font = ImageFont.truetype(font_path, 20) + + +def push_bbox(x, y, width, height, draw_boxes): + bbox = [x, y, width, height] + draw_boxes.append(bbox) + +# convert tracking box into internal box for further processing +def push_tracking(frame_index, tracking, part, dummy_box, draw_boxes): + if tracking: + width = int(tracking[4]) + height = int(tracking[5]) + x1 = max(int(tracking[0]), 0) + y1 = max(int(tracking[1]), 0) + x2 = min(int(tracking[2]), width-1) + y2 = min(int(tracking[3]), height-1) + bbox = [x1, y1, x2, y2, part] + if draw_boxes.get(frame_index) is None: + draw_boxes[frame_index] = [] + if (x1 < x2 < width + and y1 < y2 < height + and x2 - x1 < width - 1 + and y2 - y1 < height - 1): + draw_boxes[frame_index].append(bbox) + else: + if dummy_box: + draw_boxes[frame_index].append([int(frame_index), 0, 0, 1, 1, part]) + +# parse tracking data from InstanceDiffusion node +def parse_tracking(tracking, body_part, person_index, dummy_box, boxes): + for body_part_name in body_part_names: + if body_part_name == body_part or body_part == "All": + body_part_data = tracking.get(body_part_name) + + if body_part_data: + # filtering out body parts belonging to other persons + for current_person in body_part_data: + # 0 for all persons, >=1 for a specific person + if person_index == 0 or current_person == person_index - 1: + current_person_body_part_data = body_part_data.get(current_person) + if current_person_body_part_data: + # iterate through frames belonging to current person + for frame_index in current_person_body_part_data: + tracking_box = current_person_body_part_data.get(frame_index) + push_tracking(frame_index, + tracking_box, + body_part_name, + dummy_box, + boxes) + + +# Drawing frame number in the top left corner of the frame +def draw_frame_index(frame_index, draw): + text = f"{frame_index}" + text_position = (0, 0) + color = (255, 0, 0) + draw.text(text_position, text, fill=color, font=font) + +# parse schedule string +def parse_schedule(frames_schedule, schedule): + if schedule: + schedule_parts = schedule.split(",") + for schedule_part in schedule_parts: + if not ':' in schedule_part: + continue + schedule_frame_index, schedule_person_index = schedule_part.split(":") + if schedule_frame_index and schedule_person_index: + frames_schedule[schedule_frame_index.strip()] = int(schedule_person_index.strip()) diff --git a/requirements.txt b/requirements.txt index aa145a9..6446c5b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,8 @@ # Update your requirements. # numpy==1.18.1 + +pillow +comfy +torchvision +torch \ No newline at end of file