File size: 1,724 Bytes
9c98083 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | import torch
import torch.nn.functional as F
class MatchImageAndMaskSize_Dolphin:
"""Resizes an image and a mask to exactly match a reference image's
runtime pixel size. Reads reference.shape directly at execution time,
so it can't drift out of sync the way two separately-wired
GetImageSize -> width/height links can."""
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image_to_resize": ("IMAGE", {"tooltip": "Image that will be resized to match the reference."}),
"mask_to_resize": ("MASK", {"tooltip": "Mask that will be resized to match the reference."}),
"reference_image": ("IMAGE", {"tooltip": "Image whose size image_to_resize and mask_to_resize will be matched to."}),
}
}
RETURN_TYPES = ("IMAGE", "MASK")
RETURN_NAMES = ("image", "mask")
FUNCTION = "match"
CATEGORY = "image/postprocessing"
def match(self, image_to_resize, mask_to_resize, reference_image):
target_h = reference_image.shape[1]
target_w = reference_image.shape[2]
img = image_to_resize.movedim(-1, 1)
img = F.interpolate(img, size=(target_h, target_w), mode="bicubic", antialias=True)
img = img.clamp(0, 1).movedim(1, -1)
mask = mask_to_resize
if mask.ndim == 2:
mask = mask.unsqueeze(0)
mask = F.interpolate(mask.unsqueeze(1), size=(target_h, target_w), mode="nearest-exact").squeeze(1)
return (img, mask)
NODE_CLASS_MAPPINGS = {
"MatchImageAndMaskSize_Dolphin": MatchImageAndMaskSize_Dolphin,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"MatchImageAndMaskSize_Dolphin": "🐬 Match Image & Mask Size (Dolphin)",
}
|