114 lines
3.8 KiB
Python
114 lines
3.8 KiB
Python
import cv2
|
|
import torch
|
|
import numpy as np
|
|
|
|
|
|
class PasteIntoFrame:
|
|
"""
|
|
根据 mask 中央矩形尺寸缩放 content，再用 *原始 mask* 把内容贴进 frame。
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"frame_image": ("IMAGE",),
|
|
"content_image": ("IMAGE",),
|
|
"mask_image": ("IMAGE",),
|
|
"threshold": (
|
|
"FLOAT",
|
|
{"default": 0.2, "min": 0.0, "max": 1.0, "step": 0.01},
|
|
),
|
|
"connectivity": ("INT", {"default": 4, "min": 4, "max": 8}),
|
|
"keep_aspect": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("composited_image",)
|
|
FUNCTION = "run"
|
|
CATEGORY = "DrawTools"
|
|
|
|
# ---------------- 主逻辑 ---------------- #
|
|
def run(
|
|
self,
|
|
frame_image,
|
|
content_image,
|
|
mask_image,
|
|
threshold=0.2,
|
|
connectivity=4,
|
|
keep_aspect=False,
|
|
):
|
|
frame = self._tensor_to_np(frame_image) # H,W,3 或 4
|
|
content = self._tensor_to_np(content_image) # Hc,Wc,3
|
|
mask = self._tensor_to_np(
|
|
mask_image,
|
|
gray=True, # 0/255
|
|
thr=threshold,
|
|
)
|
|
|
|
H, W = mask.shape
|
|
|
|
# ① 生成 inner mask 仅用于 bbox 计算
|
|
pad = cv2.copyMakeBorder(mask, 1, 1, 1, 1, cv2.BORDER_CONSTANT, value=0)
|
|
flags = 4 if connectivity == 4 else 8
|
|
cv2.floodFill(pad, None, (0, 0), 0, flags=flags | (255 << 8))
|
|
inner = pad[1:-1, 1:-1]
|
|
|
|
ys, xs = np.where(inner == 255)
|
|
if ys.size == 0:
|
|
raise ValueError("mask 中央区域为空！")
|
|
|
|
y0, y1 = ys.min(), ys.max()
|
|
x0, x1 = xs.min(), xs.max()
|
|
target_w, target_h = x1 - x0 + 1, y1 - y0 + 1
|
|
|
|
# ② 把 content 缩放/填充到和 frame 同尺寸的 canvas
|
|
canvas = frame[:, :, :3].copy() # 复制框的 RGB (不含 alpha)
|
|
if keep_aspect:
|
|
h0, w0 = content.shape[:2]
|
|
scale = min(target_w / w0, target_h / h0)
|
|
new_w, new_h = int(w0 * scale), int(h0 * scale)
|
|
resized = cv2.resize(content, (new_w, new_h), interpolation=cv2.INTER_AREA)
|
|
y_off = y0 + (target_h - new_h) // 2
|
|
x_off = x0 + (target_w - new_w) // 2
|
|
canvas[y_off : y_off + new_h, x_off : x_off + new_w] = resized
|
|
else:
|
|
resized = cv2.resize(
|
|
content, (target_w, target_h), interpolation=cv2.INTER_AREA
|
|
)
|
|
canvas[y0 : y1 + 1, x0 : x1 + 1] = resized
|
|
|
|
# ③ 用 **原始 mask** 贴图
|
|
result = frame[:, :, :3].copy()
|
|
cv2.copyTo(canvas, mask, result) # mask=原始 => 保留全部异形区域
|
|
|
|
# ④ 恢复透明度（如果 frame 有 α）
|
|
if frame.shape[2] == 4:
|
|
alpha = frame[:, :, 3:]
|
|
result = np.concatenate([result, alpha], axis=2)
|
|
|
|
return (self._np_to_tensor(result),)
|
|
|
|
# --------- tensor ↔ numpy 工具 --------- #
|
|
def _tensor_to_np(self, img, gray=False, thr=0.2):
|
|
if isinstance(img, list):
|
|
img = img[0]
|
|
if img.ndim == 3 and img.shape[0] <= 4: # [C,H,W] → [H,W,C]
|
|
img = img.permute(1, 2, 0)
|
|
arr = img.cpu().numpy()
|
|
if arr.dtype != np.bool_ and arr.max() <= 1.0:
|
|
arr = arr * 255.0
|
|
arr = arr.astype("uint8")
|
|
if gray:
|
|
if arr.ndim == 3:
|
|
arr = cv2.cvtColor(arr, cv2.COLOR_BGR2GRAY)
|
|
arr = (arr > int(thr * 255)).astype("uint8") * 255
|
|
return arr
|
|
|
|
def _np_to_tensor(self, arr):
|
|
if arr.ndim == 2:
|
|
arr = arr[:, :, None]
|
|
tensor = torch.from_numpy(arr.astype("float32") / 255.0).permute(2, 0, 1)
|
|
return tensor.unsqueeze(0)
|