142 lines
5.2 KiB
Python
142 lines
5.2 KiB
Python
# Rect / Mask — build a MASK from a RECT, with optional feather/invert/combine
|
|
import math
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
def _image_size(image):
|
|
if isinstance(image, torch.Tensor):
|
|
if image.dim() == 4: # [B,H,W,C]
|
|
return int(image.shape[2]), int(image.shape[1])
|
|
if image.dim() == 3: # [H,W,C]
|
|
return int(image.shape[1]), int(image.shape[0])
|
|
return 512, 512
|
|
|
|
def _clamp_rect(x, y, w, h, W, H):
|
|
x = max(0, min(int(x), max(0, W - 1)))
|
|
y = max(0, min(int(y), max(0, H - 1)))
|
|
w = max(1, min(int(w), W - x))
|
|
h = max(1, min(int(h), H - y))
|
|
return x, y, w, h
|
|
|
|
def _ensure_mask_shape(mask, B, H, W, device):
|
|
# Accept [H,W], [B,H,W], or [B,1,H,1] quirky shapes from some packs
|
|
if mask is None:
|
|
return None
|
|
if mask.dim() == 2:
|
|
mask = mask.unsqueeze(0) # [1,H,W]
|
|
if mask.dim() == 4 and mask.shape[1] == 1 and mask.shape[3] == 1:
|
|
mask = mask[:, 0, :, :] # [B,H,W]
|
|
if mask.dim() != 3:
|
|
raise ValueError(f"RectMask: unsupported mask shape {tuple(mask.shape)}")
|
|
# Broadcast or clamp batch as needed
|
|
if mask.shape[0] == 1 and B > 1:
|
|
mask = mask.expand(B, H, W).clone()
|
|
elif mask.shape[0] != B:
|
|
# If sizes mismatch, just take first and broadcast
|
|
mask = mask[:1].expand(B, H, W).clone()
|
|
# Resize if spatial size mismatches
|
|
if (mask.shape[1] != H) or (mask.shape[2] != W):
|
|
mask = F.interpolate(mask.unsqueeze(1), size=(H, W), mode="bilinear", align_corners=False).squeeze(1)
|
|
return mask.to(device=device, dtype=torch.float32).clamp(0.0, 1.0)
|
|
|
|
def _gaussian_kernel1d(radius, sigma, device):
|
|
# radius: pixels; kernel size = 2*radius+1
|
|
xs = torch.arange(-radius, radius + 1, device=device, dtype=torch.float32)
|
|
k = torch.exp(-(xs**2) / (2 * sigma * sigma))
|
|
k /= k.sum().clamp_min(1e-8)
|
|
return k
|
|
|
|
def _gaussian_blur(mask, radius):
|
|
# mask: [B,H,W], radius >= 1
|
|
if radius < 1:
|
|
return mask
|
|
B, H, W = mask.shape
|
|
device = mask.device
|
|
sigma = max(0.5, radius / 2.5)
|
|
k1d = _gaussian_kernel1d(radius, sigma, device)
|
|
# separable blur: first horizontal, then vertical
|
|
x = mask.unsqueeze(1) # [B,1,H,W]
|
|
kh = k1d.view(1, 1, 1, -1)
|
|
kv = k1d.view(1, 1, -1, 1)
|
|
x = F.pad(x, (radius, radius, 0, 0), mode="reflect")
|
|
x = F.conv2d(x, kh)
|
|
x = F.pad(x, (0, 0, radius, radius), mode="reflect")
|
|
x = F.conv2d(x, kv)
|
|
return x.squeeze(1).clamp(0.0, 1.0)
|
|
|
|
class RectMask:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"rect": ("RECT",), # {"x","y","w","h"}
|
|
"feather": ("INT", {"default": 0, "min": 0, "max": 256}),
|
|
"invert": ("BOOLEAN", {"default": False}),
|
|
"combine": ("STRING", {"default": "replace",
|
|
"choices": ["replace", "union", "intersect", "subtract", "multiply"]}),
|
|
},
|
|
"optional": {
|
|
"existing_mask": ("MASK",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MASK",)
|
|
RETURN_NAMES = ("mask",)
|
|
FUNCTION = "run"
|
|
CATEGORY = "Rect"
|
|
|
|
def run(self, image, rect, feather, invert, combine, existing_mask=None):
|
|
if not isinstance(image, torch.Tensor):
|
|
raise ValueError("RectMask: expected torch.Tensor IMAGE")
|
|
|
|
# Parse rect
|
|
try:
|
|
x = int(rect.get("x", 0)); y = int(rect.get("y", 0))
|
|
w = int(rect.get("w", 1)); h = int(rect.get("h", 1))
|
|
except Exception:
|
|
x, y, w, h = 0, 0, 1, 1
|
|
|
|
# Get sizes and clamp rect
|
|
if image.dim() == 4:
|
|
B, H, W, C = int(image.shape[0]), int(image.shape[1]), int(image.shape[2]), int(image.shape[3])
|
|
elif image.dim() == 3:
|
|
B, H, W, C = 1, int(image.shape[0]), int(image.shape[1]), int(image.shape[2])
|
|
else:
|
|
raise ValueError(f"RectMask: unsupported IMAGE dims {tuple(image.shape)}")
|
|
|
|
device = image.device
|
|
x, y, w, h = _clamp_rect(x, y, w, h, W, H)
|
|
|
|
# Build binary rect mask
|
|
mask = torch.zeros((B, H, W), device=device, dtype=torch.float32)
|
|
mask[:, y:y+h, x:x+w] = 1.0
|
|
|
|
# Feather (Gaussian)
|
|
if feather > 0:
|
|
radius = int(feather)
|
|
mask = _gaussian_blur(mask, radius)
|
|
|
|
# Invert
|
|
if invert:
|
|
mask = 1.0 - mask
|
|
|
|
# Combine with existing_mask
|
|
if existing_mask is not None:
|
|
em = _ensure_mask_shape(existing_mask, B, H, W, device)
|
|
if combine == "replace":
|
|
mask = mask
|
|
elif combine == "union":
|
|
mask = torch.maximum(em, mask)
|
|
elif combine == "intersect":
|
|
mask = torch.minimum(em, mask)
|
|
elif combine == "subtract":
|
|
mask = (em - mask).clamp(0.0, 1.0)
|
|
elif combine == "multiply":
|
|
mask = (em * mask).clamp(0.0, 1.0)
|
|
|
|
return (mask,)
|
|
|
|
NODE_CLASS_MAPPINGS = {"RectMask": RectMask}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {"RectMask": "Rect / Mask"}
|