127 lines
4.9 KiB
Python
127 lines
4.9 KiB
Python
import torch
|
||
|
||
MAX_MASKS = 8
|
||
|
||
|
||
class AutoMask:
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
optional = {f"mask_{i}": ("MASK",) for i in range(1, MAX_MASKS + 1)}
|
||
optional["else"] = ("MASK",)
|
||
return {
|
||
"required": {
|
||
"image": ("IMAGE",),
|
||
"multi_match": (["Use First Match", "Raise Error"],),
|
||
"no_match": (["Output Empty Mask", "Use Else Mask", "Raise Error"],),
|
||
},
|
||
"optional": optional,
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE", "MASK", "IMAGE")
|
||
RETURN_NAMES = ("image", "mask", "rgba")
|
||
FUNCTION = "execute"
|
||
CATEGORY = "IMGLAB"
|
||
|
||
def execute(self, image, multi_match, no_match, **kwargs):
|
||
B, img_h, img_w, C = image.shape
|
||
|
||
match_results = {}
|
||
mask_sizes = {}
|
||
matched_masks = []
|
||
|
||
for i in range(1, MAX_MASKS + 1):
|
||
mask = kwargs.get(f"mask_{i}")
|
||
if mask is None:
|
||
continue
|
||
|
||
if mask.dim() == 2:
|
||
mh, mw = mask.shape
|
||
mask = mask.unsqueeze(0).expand(B, -1, -1)
|
||
elif mask.dim() == 3:
|
||
mh, mw = mask.shape[1], mask.shape[2]
|
||
if mask.shape[0] != B:
|
||
mask = mask.expand(B, -1, -1)
|
||
else:
|
||
match_results[i] = False
|
||
continue
|
||
|
||
mask_sizes[i] = [mw, mh]
|
||
is_match = (mh == img_h and mw == img_w)
|
||
match_results[i] = is_match
|
||
|
||
if is_match:
|
||
matched_masks.append((i, mask))
|
||
|
||
# ── else mask size for UI display ────────────────────────────────
|
||
else_mask_size = None
|
||
else_mask_raw = kwargs.get("else")
|
||
if else_mask_raw is not None:
|
||
if else_mask_raw.dim() == 2:
|
||
emh, emw = else_mask_raw.shape
|
||
elif else_mask_raw.dim() == 3:
|
||
emh, emw = else_mask_raw.shape[1], else_mask_raw.shape[2]
|
||
else:
|
||
emh, emw = 0, 0
|
||
else_mask_size = [emw, emh]
|
||
|
||
# ── Multiple matches ─────────────────────────────────────────────
|
||
if len(matched_masks) > 1 and multi_match == "Raise Error":
|
||
slots = [str(s) for s, _ in matched_masks]
|
||
raise ValueError(
|
||
f"[AutoPickMask] Multiple masks match {img_w}×{img_h}: "
|
||
f"slots {', '.join(slots)}. "
|
||
f"Disconnect duplicates or switch to 'Use First Match'."
|
||
)
|
||
|
||
# ── No match ─────────────────────────────────────────────────────
|
||
if len(matched_masks) == 0:
|
||
if no_match == "Raise Error":
|
||
connected = [str(k) for k in mask_sizes]
|
||
raise ValueError(
|
||
f"[AutoPickMask] No mask matches {img_w}×{img_h}. "
|
||
f"Connected slots: {', '.join(connected) or 'none'}."
|
||
)
|
||
elif no_match == "Use Else Mask":
|
||
else_mask = else_mask_raw
|
||
if else_mask is None:
|
||
raise ValueError(
|
||
f"[AutoPickMask] No mask matches {img_w}×{img_h} and "
|
||
f"'else' is not connected. Connect an else mask or "
|
||
f"change 'no_match' to 'Output Empty Mask'."
|
||
)
|
||
if else_mask.dim() == 2:
|
||
else_mask = else_mask.unsqueeze(0).expand(B, -1, -1)
|
||
elif else_mask.dim() == 3:
|
||
if else_mask.shape[0] != B:
|
||
else_mask = else_mask.expand(B, -1, -1)
|
||
matched_mask = else_mask
|
||
else:
|
||
matched_mask = torch.zeros(
|
||
B, img_h, img_w, dtype=image.dtype, device=image.device
|
||
)
|
||
else:
|
||
matched_mask = matched_masks[0][1]
|
||
|
||
# ── Resize if needed (else mask may differ in size) ──────────────
|
||
mh, mw = matched_mask.shape[1], matched_mask.shape[2]
|
||
if mh != img_h or mw != img_w:
|
||
import torch.nn.functional as F
|
||
matched_mask = F.interpolate(
|
||
matched_mask.unsqueeze(1).float(),
|
||
size=(img_h, img_w),
|
||
mode="nearest"
|
||
).squeeze(1).to(image.dtype)
|
||
|
||
alpha = matched_mask.unsqueeze(-1)
|
||
rgba = torch.cat([image, alpha], dim=-1)
|
||
|
||
return {
|
||
"result": (image, matched_mask, rgba),
|
||
"ui": {
|
||
"match_results": [match_results],
|
||
"image_size": [[img_w, img_h]],
|
||
"mask_sizes": [mask_sizes],
|
||
"else_mask_size": [else_mask_size] if else_mask_size else [],
|
||
},
|
||
}
|