| import torch |
| import torchvision.transforms.functional as TF |
| from PIL import Image, ImageDraw |
|
|
| from .Utils import get_sampler_by_name |
|
|
|
|
| class ImageTransformResizeAbsolute: |
| def __init__(self): |
| pass |
|
|
| @classmethod |
| def INPUT_TYPES(cls): |
| return { |
| "required": { |
| "images": ("IMAGE",), |
| "width": ("INT", { |
| "default": 256, |
| "min": 1, |
| "step": 1 |
| }), |
| "height": ("INT", { |
| "default": 256, |
| "min": 1, |
| "step": 1 |
| }), |
| "method": (["lanczos", "bicubic", "hamming", "bilinear", "box", "nearest"],), |
| }, |
| } |
|
|
| RETURN_TYPES = ("IMAGE",) |
| FUNCTION = "node" |
| CATEGORY = "image/transform" |
|
|
| def node(self, images, width, height, method): |
| def resize_tensor(tensor): |
| return tensor.tensor_to_image().resize((width, height), get_sampler_by_name(method)).image_to_tensor() |
|
|
| return (torch.stack([ |
| resize_tensor(images[i]) for i in range(len(images)) |
| ]),) |
|
|
|
|
| class ImageTransformResizeRelative: |
| def __init__(self): |
| pass |
|
|
| @classmethod |
| def INPUT_TYPES(cls): |
| return { |
| "required": { |
| "images": ("IMAGE",), |
| "scale_width": ("FLOAT", { |
| "default": 1.0, |
| "step": 0.1 |
| }), |
| "scale_height": ("FLOAT", { |
| "default": 1.0, |
| "step": 0.1 |
| }), |
| "method": (["lanczos", "bicubic", "hamming", "bilinear", "box", "nearest"],), |
| }, |
| } |
|
|
| RETURN_TYPES = ("IMAGE",) |
| FUNCTION = "node" |
| CATEGORY = "image/transform" |
|
|
| def node(self, images, scale_width, scale_height, method): |
| height, width = images[0, :, :, 0].shape |
|
|
| width = int(width * scale_width) |
| height = int(height * scale_height) |
|
|
| return ImageTransformResizeAbsolute().node(images, width, height, method) |
|
|
|
|
| class ImageTransformResizeClip: |
| def __init__(self): |
| pass |
|
|
| @classmethod |
| def INPUT_TYPES(cls): |
| return { |
| "required": { |
| "images": ("IMAGE",), |
| "max_width": ("INT", { |
| "default": 1024, |
| }), |
| "max_height": ("INT", { |
| "default": 1024, |
| }), |
| "min_width": ("INT", { |
| "default": 0, |
| }), |
| "min_height": ("INT", { |
| "default": 0, |
| }), |
| "method": (["lanczos", "bicubic", "hamming", "bilinear", "box", "nearest"],), |
| }, |
| } |
|
|
| RETURN_TYPES = ("IMAGE",) |
| FUNCTION = "node" |
| CATEGORY = "image/transform" |
|
|
| def node(self, images, max_width, max_height, min_width, min_height, method): |
| height, width = images[0, :, :, 0].shape |
|
|
| if min_width >= max_width or min_height >= max_height: |
| return (images,) |
|
|
| scale_min = max(min_width / width, min_height / height) |
| scale_max = min(max_width / width, max_height / height) |
|
|
| scale = max(scale_min, scale_max) |
|
|
| return ImageTransformResizeRelative().node(images, scale, scale, method) |
|
|
|
|
| class ImageTransformCropAbsolute: |
| def __init__(self): |
| pass |
|
|
| @classmethod |
| def INPUT_TYPES(cls): |
| return { |
| "required": { |
| "images": ("IMAGE",), |
| "start_x": ("INT", { |
| "default": 0, |
| "step": 1 |
| }), |
| "start_y": ("INT", { |
| "default": 0, |
| "step": 1 |
| }), |
| "end_x": ("INT", { |
| "default": 128, |
| "step": 1 |
| }), |
| "end_y": ("INT", { |
| "default": 128, |
| "step": 1 |
| }), |
| }, |
| } |
|
|
| RETURN_TYPES = ("IMAGE",) |
| FUNCTION = "node" |
| CATEGORY = "image/transform" |
|
|
| def node(self, images, start_x, start_y, end_x, end_y): |
| def resize_tensor(tensor): |
| return tensor.tensor_to_image().crop([start_x, start_y, end_x, end_y]).image_to_tensor() |
|
|
| return (torch.stack([ |
| resize_tensor(images[i]) for i in range(len(images)) |
| ]),) |
|
|
|
|
| class ImageTransformCropRelative: |
| def __init__(self): |
| pass |
|
|
| @classmethod |
| def INPUT_TYPES(cls): |
| return { |
| "required": { |
| "images": ("IMAGE",), |
| "start_x": ("FLOAT", { |
| "default": 0.25, |
| "max": 1.0, |
| "step": 0.01 |
| }), |
| "start_y": ("FLOAT", { |
| "default": 0.25, |
| "max": 1.0, |
| "step": 0.01 |
| }), |
| "end_x": ("FLOAT", { |
| "default": 0.75, |
| "max": 1.0, |
| "step": 0.01 |
| }), |
| "end_y": ("FLOAT", { |
| "default": 0.75, |
| "max": 1.0, |
| "step": 0.01 |
| }), |
| }, |
| } |
|
|
| RETURN_TYPES = ("IMAGE",) |
| FUNCTION = "node" |
| CATEGORY = "image/transform" |
|
|
| def node(self, images, start_x, start_y, end_x, end_y): |
| height, width = images[0, :, :, 0].shape |
|
|
| return ImageTransformCropAbsolute().node( |
| images, |
| width * start_x, |
| height * start_y, |
| width * end_x, |
| height * end_y |
| ) |
|
|
|
|
| class ImageTransformCropCorners: |
| def __init__(self): |
| pass |
|
|
| @classmethod |
| def INPUT_TYPES(cls): |
| return { |
| "required": { |
| "images": ("IMAGE",), |
| "radius": ("INT", { |
| "default": 180, |
| "max": 360, |
| "step": 1 |
| }), |
| "top_left_corner": (["true", "false"],), |
| "top_right_corner": (["true", "false"],), |
| "bottom_right_corner": (["true", "false"],), |
| "bottom_left_corner": (["true", "false"],), |
| "SSAA": ("INT", { |
| "default": 4, |
| "min": 1, |
| "max": 16, |
| "step": 1 |
| }), |
| "method": (["lanczos", "bicubic", "hamming", "bilinear", "box", "nearest"],), |
| }, |
| } |
|
|
| RETURN_TYPES = ("IMAGE",) |
| FUNCTION = "node" |
| CATEGORY = "image/transform" |
|
|
| |
| def node( |
| self, |
| images, |
| radius, |
| top_left_corner, |
| top_right_corner, |
| bottom_right_corner, |
| bottom_left_corner, |
| SSAA, |
| method |
| ): |
| sampler = get_sampler_by_name(method) |
|
|
| height, width = images[0, :, :, 0].shape |
|
|
| canvas = Image.new("RGBA", (width * SSAA, height * SSAA), (0, 0, 0, 0)) |
| draw = ImageDraw.Draw(canvas) |
|
|
| draw.rounded_rectangle( |
| ((0, 0), (width * SSAA, height * SSAA)), |
| radius * SSAA, (255, 255, 255, 255), |
| corners=( |
| True if top_left_corner == "true" else False, |
| True if top_right_corner == "true" else False, |
| True if bottom_right_corner == "true" else False, |
| True if bottom_left_corner == "true" else False |
| ) |
| ) |
|
|
| canvas = canvas.resize((width, height), sampler) |
| mask = 1.0 - canvas.image_to_tensor()[:, :, 3] |
|
|
| def crop_tensor(tensor): |
| return torch.stack([ |
| (tensor[:, :, i] - mask).clamp(0, 1) for i in range(tensor.shape[2]) |
| ], dim=2) |
|
|
| return (torch.stack([ |
| crop_tensor(images[i]) for i in range(len(images)) |
| ]),) |
|
|
|
|
| class ImageTransformPaddingAbsolute: |
| def __init__(self): |
| pass |
|
|
| @classmethod |
| def INPUT_TYPES(cls): |
| return { |
| "required": { |
| "images": ("IMAGE",), |
| "add_width": ("INT", { |
| "default": 64, |
| "min": 0, |
| }), |
| "add_height": ("INT", { |
| "default": 64, |
| "min": 0, |
| }), |
| "method": (["reflect", "edge", "constant"],), |
| }, |
| } |
|
|
| RETURN_TYPES = ("IMAGE",) |
| FUNCTION = "node" |
| CATEGORY = "image/transform" |
|
|
| def node(self, images, add_width, add_height, method): |
| def transpose_tensor(image): |
| tensor = image.clone().detach() |
| tensor_pad = TF.pad(tensor.permute(2, 0, 1), [add_height, add_width], padding_mode=method).permute(1, 2, 0) |
|
|
| return tensor_pad |
|
|
| return (torch.stack([ |
| transpose_tensor(images[i]) for i in range(len(images)) |
| ]),) |
|
|
|
|
| class ImageTransformPaddingRelative: |
| def __init__(self): |
| pass |
|
|
| @classmethod |
| def INPUT_TYPES(cls): |
| return { |
| "required": { |
| "images": ("IMAGE",), |
| "scale_width": ("FLOAT", { |
| "default": 0.25, |
| "step": 0.1, |
| }), |
| "scale_height": ("FLOAT", { |
| "default": 0.25, |
| "step": 0.1, |
| }), |
| "method": (["reflect", "edge", "constant"],), |
| }, |
| } |
|
|
| RETURN_TYPES = ("IMAGE",) |
| FUNCTION = "node" |
| CATEGORY = "image/transform" |
|
|
| def node(self, images, scale_width, scale_height, method): |
| height, width = images[0, :, :, 0].shape |
|
|
| add_width = int(width * scale_width) |
| add_height = int(height * scale_height) |
|
|
| return ImageTransformPaddingAbsolute().node(images, add_width, add_height, method) |
|
|
|
|
| class ImageTransformRotate: |
| def __init__(self): |
| pass |
|
|
| @classmethod |
| def INPUT_TYPES(cls): |
| return { |
| "required": { |
| "images": ("IMAGE",), |
| "angle": ("FLOAT", { |
| "default": 35.0, |
| "max": 360.0, |
| "step": 0.1 |
| }), |
| "expand": (["true", "false"],), |
| "SSAA": ("INT", { |
| "default": 4, |
| "min": 1, |
| "max": 16, |
| "step": 1 |
| }), |
| "method": (["lanczos", "bicubic", "hamming", "bilinear", "box", "nearest"],), |
| }, |
| } |
|
|
| RETURN_TYPES = ("IMAGE",) |
| FUNCTION = "node" |
| CATEGORY = "image/transform" |
|
|
| def node(self, images, angle, expand, SSAA, method): |
| height, width = images[0, :, :, 0].shape |
|
|
| def rotate_tensor(tensor): |
| if method == "lanczos": |
| resize_sampler = Image.LANCZOS |
| rotate_sampler = Image.BICUBIC |
| elif method == "bicubic": |
| resize_sampler = Image.BICUBIC |
| rotate_sampler = Image.BICUBIC |
| elif method == "hamming": |
| resize_sampler = Image.HAMMING |
| rotate_sampler = Image.BILINEAR |
| elif method == "bilinear": |
| resize_sampler = Image.BILINEAR |
| rotate_sampler = Image.BILINEAR |
| elif method == "box": |
| resize_sampler = Image.BOX |
| rotate_sampler = Image.NEAREST |
| elif method == "nearest": |
| resize_sampler = Image.NEAREST |
| rotate_sampler = Image.NEAREST |
| else: |
| raise ValueError() |
|
|
| if SSAA > 1: |
| img = tensor.tensor_to_image() |
| img_us_scaled = img.resize((width * SSAA, height * SSAA), resize_sampler) |
| img_rotated = img_us_scaled.rotate(angle, rotate_sampler, expand == "true", fillcolor=(0, 0, 0, 0)) |
| img_down_scaled = img_rotated.resize((img_rotated.width // SSAA, img_rotated.height // SSAA), resize_sampler) |
| result = img_down_scaled.image_to_tensor() |
| else: |
| img = tensor.tensor_to_image() |
| img_rotated = img.rotate(angle, rotate_sampler, expand == "true", fillcolor=(0, 0, 0, 0)) |
| result = img_rotated.image_to_tensor() |
|
|
| return result |
|
|
| if angle == 0.0 or angle == 360.0: |
| return (images,) |
| else: |
| return (torch.stack([ |
| rotate_tensor(images[i]) for i in range(len(images)) |
| ]),) |
|
|
|
|
| class ImageTransformTranspose: |
| def __init__(self): |
| pass |
|
|
| @classmethod |
| def INPUT_TYPES(cls): |
| return { |
| "required": { |
| "images": ("IMAGE",), |
| "method": (["flip_horizontally", "flip_vertically", "rotate_90", "rotate_180", "rotate_270", "transpose", "transverse"],), |
| }, |
| } |
|
|
| RETURN_TYPES = ("IMAGE",) |
| FUNCTION = "node" |
| CATEGORY = "image/transform" |
|
|
| def node(self, images, method): |
| def transpose_tensor(tensor): |
| if method == "flip_horizontally": |
| transpose = Image.FLIP_LEFT_RIGHT |
| elif method == "flip_vertically": |
| transpose = Image.FLIP_TOP_BOTTOM |
| elif method == "rotate_90": |
| transpose = Image.ROTATE_90 |
| elif method == "rotate_180": |
| transpose = Image.ROTATE_180 |
| elif method == "rotate_270": |
| transpose = Image.ROTATE_270 |
| elif method == "transpose": |
| transpose = Image.TRANSPOSE |
| elif method == "transverse": |
| transpose = Image.TRANSVERSE |
| else: |
| raise ValueError() |
|
|
| return tensor.tensor_to_image().transpose(transpose).image_to_tensor() |
|
|
| return (torch.stack([ |
| transpose_tensor(images[i]) for i in range(len(images)) |
| ]),) |
|
|
|
|
| NODE_CLASS_MAPPINGS = { |
| "ImageTransformResizeAbsolute": ImageTransformResizeAbsolute, |
| "ImageTransformResizeRelative": ImageTransformResizeRelative, |
| "ImageTransformResizeClip": ImageTransformResizeClip, |
| "ImageTransformCropAbsolute": ImageTransformCropAbsolute, |
| "ImageTransformCropRelative": ImageTransformCropRelative, |
| "ImageTransformCropCorners": ImageTransformCropCorners, |
| "ImageTransformPaddingAbsolute": ImageTransformPaddingAbsolute, |
| "ImageTransformPaddingRelative": ImageTransformPaddingRelative, |
| "ImageTransformRotate": ImageTransformRotate, |
| "ImageTransformTranspose": ImageTransformTranspose |
| } |
|
|
