Skip to content

Commit f6c5d7c

Browse files
committed
feat: add bboxes input to Create Bounding Boxes node
1 parent 7747c34 commit f6c5d7c

1 file changed

Lines changed: 133 additions & 3 deletions

File tree

comfy_extras/nodes_bounding_boxes.py

Lines changed: 133 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
import json
2+
13
import numpy as np
24
import torch
35
from PIL import Image, ImageDraw, ImageEnhance, ImageFont
@@ -166,6 +168,111 @@ def boxes_to_regions(boxes, width: int, height: int) -> list:
166168
return regions
167169

168170

171+
def normalize_incoming_boxes(bboxes) -> list:
172+
if isinstance(bboxes, dict):
173+
frame = [bboxes]
174+
elif not isinstance(bboxes, list) or not bboxes:
175+
frame = []
176+
elif isinstance(bboxes[0], dict):
177+
frame = bboxes
178+
else:
179+
frame = bboxes[0] if isinstance(bboxes[0], list) else []
180+
boxes = []
181+
for box in frame:
182+
if not isinstance(box, dict):
183+
continue
184+
norm = {
185+
"x": box.get("x", 0),
186+
"y": box.get("y", 0),
187+
"width": box.get("width", 0),
188+
"height": box.get("height", 0),
189+
}
190+
meta = box.get("metadata")
191+
if isinstance(meta, dict):
192+
norm["metadata"] = meta
193+
boxes.append(norm)
194+
return boxes
195+
196+
197+
def _looks_like_element(box: dict) -> bool:
198+
bbox = box.get("bbox")
199+
return isinstance(bbox, (list, tuple)) and len(bbox) == 4
200+
201+
202+
def _looks_like_bbox(box: dict) -> bool:
203+
return all(key in box for key in ("x", "y", "width", "height"))
204+
205+
206+
def elements_to_boxes(elements: list, width: int, height: int) -> list:
207+
boxes = []
208+
for element in elements:
209+
if not isinstance(element, dict):
210+
continue
211+
bbox = element.get("bbox")
212+
if not (isinstance(bbox, (list, tuple)) and len(bbox) == 4):
213+
raise ValueError("bboxes element is missing a valid 'bbox' [ymin, xmin, ymax, xmax]")
214+
try:
215+
ymin, xmin, ymax, xmax = (float(v) / 1000.0 for v in bbox)
216+
except (TypeError, ValueError):
217+
raise ValueError("bboxes element 'bbox' must contain four numbers")
218+
etype = "text" if element.get("type") == "text" else "obj"
219+
boxes.append({
220+
"x": round(min(xmin, xmax) * width),
221+
"y": round(min(ymin, ymax) * height),
222+
"width": round(abs(xmax - xmin) * width),
223+
"height": round(abs(ymax - ymin) * height),
224+
"metadata": {
225+
"type": etype,
226+
"text": element.get("text", "") if etype == "text" else "",
227+
"desc": element.get("desc", ""),
228+
"palette": element.get("color_palette", []) or [],
229+
},
230+
})
231+
return boxes
232+
233+
234+
def boxes_from_input(data, width: int, height: int) -> list:
235+
if data is None:
236+
return []
237+
if isinstance(data, str):
238+
text = data.strip()
239+
if not text:
240+
return []
241+
try:
242+
data = json.loads(text)
243+
except (ValueError, TypeError) as exc:
244+
raise ValueError(f"bboxes string input is not valid JSON: {exc}") from exc
245+
if isinstance(data, dict):
246+
if _looks_like_element(data):
247+
return elements_to_boxes([data], width, height)
248+
if _looks_like_bbox(data):
249+
return normalize_incoming_boxes(data)
250+
raise ValueError(
251+
"bboxes dict must be a bounding box (x, y, width, height) or an element (with a 'bbox')"
252+
)
253+
if not isinstance(data, list):
254+
raise ValueError(
255+
"bboxes input must be bounding boxes, elements, or a JSON string, "
256+
f"got {type(data).__name__}"
257+
)
258+
if not data:
259+
return []
260+
first = data[0]
261+
if isinstance(first, list):
262+
return normalize_incoming_boxes(data)
263+
if isinstance(first, dict):
264+
if _looks_like_element(first):
265+
return elements_to_boxes(data, width, height)
266+
if _looks_like_bbox(first):
267+
return normalize_incoming_boxes(data)
268+
raise ValueError(
269+
"bboxes items must be bounding boxes (x, y, width, height) or elements (with a 'bbox')"
270+
)
271+
raise ValueError(
272+
f"bboxes list must contain bounding boxes or elements, got {type(first).__name__}"
273+
)
274+
275+
169276
def _norm_bbox(region: dict) -> list[int]:
170277
def grid(value: float) -> int:
171278
return max(0, min(1000, round(value * 1000)))
@@ -199,6 +306,8 @@ def build_elements(regions: list) -> list:
199306

200307

201308
class CreateBoundingBoxes(io.ComfyNode):
309+
_last_incoming: dict = {}
310+
202311
@classmethod
203312
def define_schema(cls):
204313
editor_state = io.BoundingBoxes.Input(
@@ -217,6 +326,12 @@ def define_schema(cls):
217326
optional=True,
218327
tooltip="Optional image used as background in the canvas and preview.",
219328
),
329+
io.MultiType.Input(
330+
"bboxes",
331+
[io.BoundingBox, io.Array, io.String],
332+
optional=True,
333+
tooltip="Bounding boxes, elements, or a JSON string to seed the canvas. A new upstream value seeds the canvas; edits you make on the canvas take priority and are kept until the upstream value changes again.",
334+
),
220335
io.Int.Input("width", default=1024, min=64, max=16384, step=16,
221336
tooltip="Width of the canvas and the pixel grid for the bounding boxes."),
222337
io.Int.Input("height", default=1024, min=64, max=16384, step=16,
@@ -228,18 +343,33 @@ def define_schema(cls):
228343
io.BoundingBox.Output(display_name="bboxes"),
229344
io.Array.Output(display_name="elements"),
230345
],
346+
hidden=[io.Hidden.unique_id],
347+
is_output_node=True,
231348
is_experimental=True,
232349
)
233350

234351
@classmethod
235-
def execute(cls, width, height, editor_state=None, background=None) -> io.NodeOutput:
236-
regions = boxes_to_regions(editor_state, width, height)
352+
def execute(cls, width, height, editor_state=None, background=None, bboxes=None) -> io.NodeOutput:
353+
incoming = boxes_from_input(bboxes, width, height)
354+
node_id = cls.hidden.unique_id
355+
if incoming:
356+
changed = cls._last_incoming.get(node_id) != incoming
357+
if changed:
358+
cls._last_incoming[node_id] = incoming
359+
else:
360+
changed = False
361+
cls._last_incoming.pop(node_id, None)
362+
source = incoming if changed else (editor_state or incoming)
363+
regions = boxes_to_regions(source, width, height)
237364
preview = render_preview(regions, width, height, _bg_from_image(background))
365+
ui = {"dims": [width, height]}
366+
if incoming:
367+
ui["input_bboxes"] = incoming
238368
return io.NodeOutput(
239369
preview,
240370
fractions_to_bbox_frame(regions, width, height),
241371
build_elements(regions),
242-
ui={"dims": [width, height]},
372+
ui=ui,
243373
)
244374

245375

0 commit comments

Comments
 (0)