import torch from PIL import Image from .imagefunc import AnyType, log, tensor2pil, pil2tensor, image2mask, fit_resize_image any = AnyType("*") class ImageMaskScaleAs: def __init__(self): self.NODE_NAME = 'ImageMaskScaleAs' @classmethod def INPUT_TYPES(self): fit_mode = ['letterbox', 'crop', 'fill'] method_mode = ['lanczos', 'bicubic', 'hamming', 'bilinear', 'box', 'nearest'] return { "required": { "scale_as": (any, {}), "fit": (fit_mode,), "method": (method_mode,), }, "optional": { "image": ("IMAGE",), # "mask": ("MASK",), # } } RETURN_TYPES = ("IMAGE", "MASK", "BOX", "INT", "INT") RETURN_NAMES = ("image", "mask", "original_size", "widht", "height",) FUNCTION = 'image_mask_scale_as' CATEGORY = '😺dzNodes/LayerUtility' def image_mask_scale_as(self, scale_as, fit, method, image=None, mask = None, ): if scale_as.shape[0] > 0: _asimage = tensor2pil(scale_as[0]) else: _asimage = tensor2pil(scale_as) target_width, target_height = _asimage.size _mask = Image.new('L', size=_asimage.size, color='black') _image = Image.new('RGB', size=_asimage.size, color='black') orig_width = 4 orig_height = 4 resize_sampler = Image.LANCZOS if method == "bicubic": resize_sampler = Image.BICUBIC elif method == "hamming": resize_sampler = Image.HAMMING elif method == "bilinear": resize_sampler = Image.BILINEAR elif method == "box": resize_sampler = Image.BOX elif method == "nearest": resize_sampler = Image.NEAREST ret_images = [] ret_masks = [] if image is not None: for i in image: i = torch.unsqueeze(i, 0) _image = tensor2pil(i).convert('RGB') orig_width, orig_height = _image.size _image = fit_resize_image(_image, target_width, target_height, fit, resize_sampler) ret_images.append(pil2tensor(_image)) if mask is not None: if mask.dim() == 2: mask = torch.unsqueeze(mask, 0) for m in mask: m = torch.unsqueeze(m, 0) _mask = tensor2pil(m).convert('L') orig_width, orig_height = _mask.size _mask = fit_resize_image(_mask, target_width, target_height, fit, resize_sampler).convert('L') ret_masks.append(image2mask(_mask)) if len(ret_images) > 0 and len(ret_masks) >0: log(f"{self.NODE_NAME} Processed {len(ret_images)} image(s).", message_type='finish') return (torch.cat(ret_images, dim=0), torch.cat(ret_masks, dim=0), [orig_width, orig_height],target_width, target_height,) elif len(ret_images) > 0 and len(ret_masks) == 0: log(f"{self.NODE_NAME} Processed {len(ret_images)} image(s).", message_type='finish') return (torch.cat(ret_images, dim=0), None, [orig_width, orig_height],target_width, target_height,) elif len(ret_images) == 0 and len(ret_masks) > 0: log(f"{self.NODE_NAME} Processed {len(ret_masks)} image(s).", message_type='finish') return (None, torch.cat(ret_masks, dim=0), [orig_width, orig_height], target_width, target_height,) else: log(f"Error: {self.NODE_NAME} skipped, because the available image or mask is not found.", message_type='error') return (None, None, [orig_width, orig_height], 0, 0,) class LS_ImageMaskScaleAsV2: def __init__(self): self.NODE_NAME = 'ImageMaskScaleAsV2' @classmethod def INPUT_TYPES(self): fit_mode = ['letterbox', 'crop', 'fill'] method_mode = ['lanczos', 'bicubic', 'hamming', 'bilinear', 'box', 'nearest'] return { "required": { "scale_as": (any, {}), "fit": (fit_mode,), "method": (method_mode,), "background_color": ("STRING", {"default": "#FFFFFF"},), }, "optional": { "image": ("IMAGE",), # "mask": ("MASK",), # } } RETURN_TYPES = ("IMAGE", "MASK", "BOX", "INT", "INT") RETURN_NAMES = ("image", "mask", "original_size", "widht", "height",) FUNCTION = 'image_mask_scale_as_v2' CATEGORY = '😺dzNodes/LayerUtility' def image_mask_scale_as_v2(self, scale_as, fit, method, background_color, image=None, mask=None, ): if scale_as.shape[0] > 0: _asimage = tensor2pil(scale_as[0]) else: _asimage = tensor2pil(scale_as) target_width, target_height = _asimage.size _mask = Image.new('L', size=_asimage.size, color='black') _image = Image.new('RGB', size=_asimage.size, color=background_color) orig_width = 4 orig_height = 4 resize_sampler = Image.LANCZOS if method == "bicubic": resize_sampler = Image.BICUBIC elif method == "hamming": resize_sampler = Image.HAMMING elif method == "bilinear": resize_sampler = Image.BILINEAR elif method == "box": resize_sampler = Image.BOX elif method == "nearest": resize_sampler = Image.NEAREST ret_images = [] ret_masks = [] if image is not None: for i in image: i = torch.unsqueeze(i, 0) _image = tensor2pil(i).convert('RGB') orig_width, orig_height = _image.size _image = fit_resize_image(_image, target_width, target_height, fit, resize_sampler, background_color=background_color) ret_images.append(pil2tensor(_image)) if mask is not None: if mask.dim() == 2: mask = torch.unsqueeze(mask, 0) for m in mask: m = torch.unsqueeze(m, 0) _mask = tensor2pil(m).convert('L') orig_width, orig_height = _mask.size _mask = fit_resize_image(_mask, target_width, target_height, fit, resize_sampler, background_color=background_color).convert('L') ret_masks.append(image2mask(_mask)) if len(ret_images) > 0 and len(ret_masks) > 0: log(f"{self.NODE_NAME} Processed {len(ret_images)} image(s).", message_type='finish') return (torch.cat(ret_images, dim=0), torch.cat(ret_masks, dim=0), [orig_width, orig_height], target_width, target_height,) elif len(ret_images) > 0 and len(ret_masks) == 0: log(f"{self.NODE_NAME} Processed {len(ret_images)} image(s).", message_type='finish') return (torch.cat(ret_images, dim=0), None, [orig_width, orig_height], target_width, target_height,) elif len(ret_images) == 0 and len(ret_masks) > 0: log(f"{self.NODE_NAME} Processed {len(ret_masks)} image(s).", message_type='finish') return (None, torch.cat(ret_masks, dim=0), [orig_width, orig_height], target_width, target_height,) else: log(f"Error: {self.NODE_NAME} skipped, because the available image or mask is not found.", message_type='error') return (None, None, [orig_width, orig_height], 0, 0,) NODE_CLASS_MAPPINGS = { "LayerUtility: ImageMaskScaleAs": ImageMaskScaleAs, "LayerUtility: ImageMaskScaleAsV2": LS_ImageMaskScaleAsV2, } NODE_DISPLAY_NAME_MAPPINGS = { "LayerUtility: ImageMaskScaleAs": "LayerUtility: Image Mask Scale As", "LayerUtility: ImageMaskScaleAsV2": "LayerUtility: Image Mask Scale As V2", }