diff --git a/comfy_extras/nodes_bounding_boxes.py b/comfy_extras/nodes_bounding_boxes.py index de3709b9100..186e1c36706 100644 --- a/comfy_extras/nodes_bounding_boxes.py +++ b/comfy_extras/nodes_bounding_boxes.py @@ -5,6 +5,7 @@ from PIL import Image, ImageDraw, ImageEnhance, ImageFont from typing_extensions import override +import nodes from comfy_api.latest import ComfyExtension, io from comfy_extras.color_util import hex_to_rgb, normalize_palette, readable_color @@ -361,6 +362,9 @@ def execute(cls, width, height, editor_state=None, last_incoming=None, backgroun ui = {"dims": [width, height]} if incoming: ui["input_bboxes"] = incoming + if background is not None and len(background) > 0: + saved = nodes.PreviewImage().save_images(background[:1], "comfy.bboxes.background") + ui["background_images"] = saved["ui"]["images"] return io.NodeOutput( preview, fractions_to_bbox_frame(regions, width, height), diff --git a/tests-unit/comfy_extras_test/nodes_bounding_boxes_test.py b/tests-unit/comfy_extras_test/nodes_bounding_boxes_test.py new file mode 100644 index 00000000000..300118484ff --- /dev/null +++ b/tests-unit/comfy_extras_test/nodes_bounding_boxes_test.py @@ -0,0 +1,40 @@ +import os + +import torch +from PIL import Image + +import folder_paths +from comfy.cli_args import args as cli_args + +if not torch.cuda.is_available(): + cli_args.cpu = True + +from comfy_extras.nodes_bounding_boxes import CreateBoundingBoxes # noqa: E402 + + +def run(background=None): + return CreateBoundingBoxes.execute(width=64, height=64, background=background).ui + + +def test_echoes_the_first_background_image_as_a_temp_preview(tmp_path, monkeypatch): + monkeypatch.setattr(folder_paths, "get_temp_directory", lambda: str(tmp_path)) + first = torch.zeros(1, 48, 80, 3) + second = torch.ones(1, 48, 80, 3) + + ui = run(background=torch.cat([first, second])) + + [entry] = ui["background_images"] + assert entry["type"] == "temp" + path = os.path.join(str(tmp_path), entry["subfolder"], entry["filename"]) + with Image.open(path) as saved: + assert saved.size == (80, 48) + assert saved.getpixel((0, 0)) == (0, 0, 0) + + +def test_omits_the_background_preview_without_a_background(tmp_path, monkeypatch): + monkeypatch.setattr(folder_paths, "get_temp_directory", lambda: str(tmp_path)) + + ui = run() + + assert "background_images" not in ui + assert ui["dims"] == [64, 64]