| 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)", |
| } |
|
|
