Added schedulers, updated to V3 schema, refactored code
This commit is contained in:
Yury Krychala
@@ -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`
|
||||
|
||||
|
||||
|
||||
+7
-18
@@ -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"]
|
||||
|
||||
@@ -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, )
|
||||
+179
-76
@@ -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)
|
||||
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())
|
||||
|
||||
@@ -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)
|
||||
@@ -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())
|
||||
@@ -1,3 +1,8 @@
|
||||
# Update your requirements.
|
||||
# numpy==1.18.1
|
||||
|
||||
|
||||
pillow
|
||||
comfy
|
||||
torchvision
|
||||
torch
|
||||
Reference in New Issue
Block a user