70 lines
2.2 KiB
Python
70 lines
2.2 KiB
Python
# "Design-focused node suite for ComfyUI."
|
|
# Copyright 2026 Augment Studio
|
|
# Augmentstudio.app
|
|
|
|
import numpy as np
|
|
import torch
|
|
import json
|
|
|
|
|
|
class MaskBoundingBox:
|
|
"""Finds the bounding box of non-zero mask content and outputs the cropped mask + coordinates."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"mask": ("MASK",),
|
|
"padding": ("INT", {"default": 0, "min": 0, "max": 512, "step": 1}),
|
|
"invert": ("BOOLEAN", {"default": False}),
|
|
"threshold": ("FLOAT", {"default": 0.10, "min": 0.01, "max": 1.0, "step": 0.01}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MASK", "INT", "INT", "INT", "INT", "AUGMENT_JSON", "TRIGGER")
|
|
RETURN_NAMES = ("mask", "x", "y", "width", "height", "json_result", "trigger")
|
|
FUNCTION = "find_bounds"
|
|
CATEGORY = "Augment/Image"
|
|
|
|
def find_bounds(self, mask, padding, invert, threshold):
|
|
m = mask[0].cpu().numpy()
|
|
img_h, img_w = m.shape
|
|
|
|
if invert:
|
|
m = 1.0 - m
|
|
|
|
rows = np.any(m > threshold, axis=1)
|
|
cols = np.any(m > threshold, axis=0)
|
|
|
|
if not rows.any():
|
|
x_min, y_min, out_w, out_h = 0, 0, img_w, img_h
|
|
else:
|
|
row_indices = np.where(rows)[0]
|
|
col_indices = np.where(cols)[0]
|
|
y_min, y_max = int(row_indices[0]), int(row_indices[-1])
|
|
x_min, x_max = int(col_indices[0]), int(col_indices[-1])
|
|
|
|
x_min = max(x_min - padding, 0)
|
|
y_min = max(y_min - padding, 0)
|
|
x_max = min(x_max + padding, img_w - 1)
|
|
y_max = min(y_max + padding, img_h - 1)
|
|
|
|
out_w = x_max - x_min + 1
|
|
out_h = y_max - y_min + 1
|
|
|
|
cropped_mask = mask[:, y_min:y_min + out_h, x_min:x_min + out_w]
|
|
|
|
print(f"[augment] bounds: x={x_min} y={y_min} w={out_w} h={out_h}")
|
|
|
|
json_result = json.dumps({
|
|
"node": "MaskBoundingBox",
|
|
"x": x_min,
|
|
"y": y_min,
|
|
"width": out_w,
|
|
"height": out_h,
|
|
"padding": padding,
|
|
"invert": invert,
|
|
})
|
|
|
|
return (cropped_mask, x_min, y_min, out_w, out_h, json_result, "done")
|