486 lines
14 KiB
Python
486 lines
14 KiB
Python
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"
|
|
|
|
# noinspection PyUnresolvedReferences, PyArgumentList
|
|
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
|
|
}
|