49 lines
1.4 KiB
Python
49 lines
1.4 KiB
Python
import cv2
|
|
import torch
|
|
import numpy as np
|
|
|
|
|
|
class DetectInnerBox:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"mask": ("MASK",), # 只接 MASK
|
|
"connectivity": ("INT", {"default": 4, "min": 4, "max": 8}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("INT", "INT") # width, height
|
|
RETURN_NAMES = ("inner_w", "inner_h")
|
|
FUNCTION = "run"
|
|
CATEGORY = "DrawTools"
|
|
|
|
# ---------------- 主逻辑 ---------------- #
|
|
def run(self, mask, connectivity=4):
|
|
# ① tensor → numpy 单通道 0/255
|
|
if isinstance(mask, list): # ComfyUI 有时包在 list 里
|
|
mask = mask[0]
|
|
if mask.ndim == 3: # [1,H,W] → [H,W]
|
|
mask = mask[0]
|
|
np_mask = mask.cpu().numpy().astype("uint8") * 255 # 0/255
|
|
|
|
H, W = np_mask.shape[:2]
|
|
|
|
# ② flood-fill 把连边的白色清零
|
|
pad = cv2.copyMakeBorder(np_mask, 1, 1, 1, 1, cv2.BORDER_CONSTANT, value=0)
|
|
flags = 4 if connectivity == 4 else 8
|
|
cv2.floodFill(
|
|
pad, None, (0, 0), 0, loDiff=0, upDiff=0, flags=flags | (255 << 8)
|
|
)
|
|
inner = pad[1:-1, 1:-1]
|
|
|
|
# ③ 取剩余像素的 bbox
|
|
ys, xs = np.where(inner == 255)
|
|
if ys.size == 0:
|
|
return 0, 0
|
|
|
|
y0, y1 = ys.min(), ys.max()
|
|
x0, x1 = xs.min(), xs.max()
|
|
|
|
return int(x1 - x0 + 1), int(y1 - y0 + 1)
|