ausboss
3379df7e7f
The mask is white inside every box and black elsewhere, at the size of the picture on background (or the canvas without one). Draw a box over the words on a sign, wire the mask into an inpaint chain, and only the words change. It comes after the other three outputs, so saved workflows keep their wires.
133 lines
6.2 KiB
Python
133 lines
6.2 KiB
Python
"""Layout Boxes 🆎: the node itself (preview, the wire's boxes, the background
|
|
for the face). Needs torch and Pillow: run it with ComfyUI's Python."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
import json
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
import torch
|
|
|
|
ROOT = Path(__file__).resolve().parent.parent
|
|
sys.path.insert(0, str(ROOT))
|
|
if "nodes" in sys.modules and not hasattr(sys.modules["nodes"], "__path__"):
|
|
del sys.modules["nodes"]
|
|
|
|
from nodes import node_layout_boxes
|
|
from nodes.node_layout_boxes import UI_KEY, AusBossLayoutBoxes
|
|
|
|
W, H = 1344, 768
|
|
|
|
|
|
def layout(*boxes, from_input=None):
|
|
data = {"v": 1, "boxes": list(boxes)}
|
|
if from_input:
|
|
data["from_input"] = from_input
|
|
return json.dumps(data)
|
|
|
|
|
|
def box(x, y, w, h, kind="obj", text="", desc="", palette=()):
|
|
return {"x": x / W, "y": y / H, "w": w / W, "h": h / H, "type": kind, "text": text, "desc": desc, "palette": list(palette)}
|
|
|
|
|
|
def fake_folder_paths(temp_dir: str):
|
|
class FakeFolderPaths:
|
|
@staticmethod
|
|
def get_temp_directory():
|
|
return temp_dir
|
|
|
|
@staticmethod
|
|
def get_save_image_path(prefix, output_dir, _width, _height):
|
|
return str(output_dir), prefix, 1, "", prefix
|
|
|
|
return FakeFolderPaths
|
|
|
|
|
|
class LayoutBoxesNode(unittest.TestCase):
|
|
def setUp(self):
|
|
self.node = AusBossLayoutBoxes()
|
|
|
|
def test_the_canvas_boxes_make_all_four_outputs(self):
|
|
out = self.node.draw(width=W, height=H, layout=layout(box(320, 96, 512, 640, desc="a cat"), box(0, 0, W, 64, "text", "OPEN", "neon")))
|
|
preview, frames, elements, mask = out["result"]
|
|
self.assertEqual(tuple(preview.shape), (1, 585, 1024, 3), "the canvas's shape, 1024 on its long side")
|
|
self.assertEqual(preview.dtype, torch.float32)
|
|
self.assertTrue(0.0 <= float(preview.min()) and float(preview.max()) <= 1.0)
|
|
self.assertEqual(frames, [[{"x": 320, "y": 96, "width": 512, "height": 640}, {"x": 0, "y": 0, "width": W, "height": 64}]])
|
|
self.assertEqual([e["type"] for e in elements], ["obj", "text"])
|
|
self.assertEqual(out["ui"][UI_KEY], [{"size": [W, H], "count": 2}])
|
|
self.assertEqual(tuple(mask.shape), (1, H, W), "the canvas size")
|
|
self.assertEqual(mask.dtype, torch.float32)
|
|
expected = torch.zeros(1, H, W)
|
|
expected[0, 96:736, 320:832] = 1.0
|
|
expected[0, 0:64, :] = 1.0
|
|
self.assertTrue(torch.equal(mask, expected), "1 inside every box, 0 elsewhere")
|
|
|
|
def test_no_boxes(self):
|
|
preview, frames, elements, mask = self.node.draw(width=512, height=1024)["result"]
|
|
self.assertEqual(tuple(preview.shape), (1, 1024, 512, 3))
|
|
self.assertEqual((frames, elements), ([], []))
|
|
self.assertEqual(tuple(mask.shape), (1, 1024, 512))
|
|
self.assertEqual(float(mask.abs().sum()), 0.0, "no boxes: all black")
|
|
|
|
def test_the_mask_fits_the_background_picture(self):
|
|
# A 1000 x 750 photo: the canvas is 1008 x 752 (16 px steps), and the
|
|
# mask has the photo's own size, the box at the same place on it.
|
|
picture = torch.rand(2, 750, 1000, 3)
|
|
words = {"x": 0.5, "y": 0.25, "w": 0.25, "h": 0.5, "type": "text", "text": "SALE", "desc": "red letters"}
|
|
mask = self.node.draw(background=picture, width=1008, height=752, layout=layout(words))["result"][3]
|
|
self.assertEqual(tuple(mask.shape), (1, 750, 1000), "one mask for the batch, at the picture's size")
|
|
self.assertEqual(int(mask.sum()), 250 * 375)
|
|
self.assertEqual(float(mask[0, 188, 500]), 1.0)
|
|
self.assertEqual(float(mask[0, 562, 749]), 1.0)
|
|
self.assertEqual(float(mask[0, 187, 500]) + float(mask[0, 563, 749]) + float(mask[0, 300, 750]), 0.0)
|
|
|
|
def test_new_boxes_on_the_wire_win_once(self):
|
|
wire = [{"type": "obj", "bbox": [0, 0, 500, 500], "desc": "from the wire"}]
|
|
drawn = layout(box(640, 384, 128, 128, desc="drawn"))
|
|
first = self.node.draw(width=W, height=H, layout=drawn, bboxes=wire)
|
|
self.assertEqual(first["result"][2][0]["desc"], "from the wire")
|
|
incoming = first["ui"][UI_KEY][0]["incoming"]
|
|
self.assertEqual(incoming["boxes"][0]["desc"], "from the wire")
|
|
# The face put them on the canvas and kept the fingerprint; a change
|
|
# made there is used from then on, while the wire stays the same.
|
|
edited = layout(dict(incoming["boxes"][0], desc="changed on the canvas"), from_input=incoming["fingerprint"])
|
|
again = self.node.draw(width=W, height=H, layout=edited, bboxes=wire)
|
|
self.assertEqual(again["result"][2][0]["desc"], "changed on the canvas")
|
|
other = self.node.draw(width=W, height=H, layout=edited, bboxes=[dict(wire[0], desc="new on the wire")])
|
|
self.assertEqual(other["result"][2][0]["desc"], "new on the wire")
|
|
|
|
def test_the_background_goes_to_the_temp_folder_for_the_face(self):
|
|
picture = torch.rand(2, 900, 600, 3)
|
|
with tempfile.TemporaryDirectory() as tmp, patch.object(node_layout_boxes, "folder_paths", fake_folder_paths(tmp)):
|
|
out = self.node.draw(background=picture, width=608, height=896, layout=layout())
|
|
note = out["ui"][UI_KEY][0]
|
|
self.assertEqual(note["picture"], [600, 900])
|
|
self.assertEqual(note["background"]["type"], "temp")
|
|
saved = Path(tmp) / note["background"]["filename"]
|
|
self.assertTrue(saved.is_file())
|
|
self.assertEqual(tuple(out["result"][0].shape), (1, 896, 608, 3), "a small canvas is not blown up")
|
|
|
|
def test_a_grey_background_and_a_failed_save_still_run(self):
|
|
picture = torch.rand(1, 64, 96, 1)
|
|
out = self.node.draw(background=picture, width=96, height=64)
|
|
self.assertNotIn("background", out["ui"][UI_KEY][0], "no folder_paths offline: no file, no error")
|
|
self.assertEqual(tuple(out["result"][0].shape), (1, 64, 96, 3))
|
|
|
|
class Broken:
|
|
@staticmethod
|
|
def get_temp_directory():
|
|
raise OSError("disk full")
|
|
|
|
with patch.object(node_layout_boxes, "folder_paths", Broken):
|
|
out = self.node.draw(background=torch.rand(1, 64, 64, 3), width=64, height=64)
|
|
self.assertNotIn("background", out["ui"][UI_KEY][0])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|