213 lines
7.0 KiB
Python
213 lines
7.0 KiB
Python
from typing import Literal, Any
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torchvision.transforms as t
|
|
from PIL import Image as ImageF
|
|
from PIL.Image import Image as ImageB
|
|
from torch import Tensor, dtype
|
|
|
|
|
|
def tensor_to_image(self):
|
|
return t.ToPILImage()(self.permute(2, 0, 1))
|
|
|
|
|
|
def image_to_tensor(self):
|
|
return t.ToTensor()(self).permute(1, 2, 0)
|
|
|
|
|
|
def create_rgba_image(width: int, height: int, color=(0, 0, 0, 0)) -> ImageB:
|
|
return ImageF.new("RGBA", (width, height), color)
|
|
|
|
|
|
def get_sampler_by_name(method) -> Literal[0, 1, 2, 3, 4, 5]:
|
|
if method == "lanczos":
|
|
return ImageF.LANCZOS
|
|
elif method == "bicubic":
|
|
return ImageF.BICUBIC
|
|
elif method == "hamming":
|
|
return ImageF.HAMMING
|
|
elif method == "bilinear":
|
|
return ImageF.BILINEAR
|
|
elif method == "box":
|
|
return ImageF.BOX
|
|
elif method == "nearest":
|
|
return ImageF.NEAREST
|
|
else:
|
|
raise ValueError("Sampler not found.")
|
|
|
|
|
|
def cv2_layer(tensor: Tensor, function) -> Tensor:
|
|
"""
|
|
This function applies a given function to each channel of an input tensor and returns the result as a PyTorch tensor.
|
|
|
|
:param tensor: A PyTorch tensor of shape (H, W, C) or (N, H, W, C), where C is the number of channels, H is the height, and W is the width of the image.
|
|
:param function: A function that takes a numpy array of shape (H, W, C) as input and returns a numpy array of the same shape.
|
|
:return: A PyTorch tensor of the same shape as the input tensor, where the given function has been applied to each channel of each image in the tensor.
|
|
"""
|
|
shape_size = tensor.shape.__len__()
|
|
|
|
def produce(image):
|
|
channels = image[0, 0, :].shape[0]
|
|
|
|
rgb = image[:, :, 0:3].numpy()
|
|
result_rgb = function(rgb)
|
|
|
|
if channels <= 3:
|
|
return torch.from_numpy(result_rgb)
|
|
elif channels == 4:
|
|
alpha = image[:, :, 3:4].numpy()
|
|
result_alpha = function(alpha)[..., np.newaxis]
|
|
result_rgba = np.concatenate((result_rgb, result_alpha), axis=2)
|
|
|
|
return torch.from_numpy(result_rgba)
|
|
|
|
if shape_size == 3:
|
|
return torch.from_numpy(produce(tensor))
|
|
elif shape_size == 4:
|
|
return torch.stack([
|
|
produce(tensor[i]) for i in range(len(tensor))
|
|
])
|
|
else:
|
|
raise ValueError("Incompatible tensor dimension.")
|
|
|
|
|
|
def radialspace_1D(
|
|
size: int,
|
|
curvy: float = 1.0,
|
|
scale: float = 1.0,
|
|
min_val: float = 0.0,
|
|
max_val: float = 1.0,
|
|
normalize: bool = True,
|
|
dtype: dtype = torch.float32
|
|
):
|
|
"""
|
|
Create a 1D tensor with a radial gradient from 0 to 1 from the center to the edges.
|
|
|
|
:param size: Tuple of int
|
|
Size of the tensor.
|
|
Can be a tuple of one or two integers.
|
|
If one integer is given, the tensor will be square.
|
|
If two integers are given, the tensor will have the given width and height.
|
|
:param curvy: Float
|
|
Curviness of the gradient.
|
|
Higher values result in a sharper transition from 0 to 1.
|
|
:param scale: Float
|
|
Scale of the gradient.
|
|
Higher values result in a larger area with values close to 1.
|
|
:param min_val: Float
|
|
Minimum value in the tensor.
|
|
:param max_val: Float
|
|
Maximum value in the tensor.
|
|
:param normalize: Bool
|
|
Whether to normalize the tensor values to be between min_val and max_val.
|
|
:param dtype: Torch.dtype
|
|
Data type of the resulting tensor.
|
|
|
|
:return: Torch.Tensor
|
|
A 2D tensor with a radial gradient from 0 to 1 from the center to the edges.
|
|
"""
|
|
|
|
tensor = radialspace_2D((1, size), curvy, scale, "square", min_val, max_val, (0.0, 0.5), None, normalize, dtype).squeeze()
|
|
|
|
if min_val < 0:
|
|
tensor[:len(tensor) // 2] = -(tensor[:len(tensor) // 2])
|
|
|
|
if normalize:
|
|
tensor = ((tensor - tensor.min()) / (tensor.max() - tensor.min())) * (max_val - min_val) + min_val
|
|
|
|
return tensor
|
|
|
|
|
|
def radialspace_2D(
|
|
size: tuple[int] | tuple[int, int],
|
|
curvy: float = 1.0,
|
|
scale: float = 1.0,
|
|
mode: str = "square",
|
|
min_val: float = 0.0,
|
|
max_val: float = 1.0,
|
|
center: tuple[float, float] = (0.5, 0.5),
|
|
function: Any = None,
|
|
normalize: bool = True,
|
|
dtype: dtype = torch.float32
|
|
) -> Tensor:
|
|
"""
|
|
Create a 2D tensor with a radial gradient from 0 to 1 from the center to the edges.
|
|
|
|
:param size: Tuple of int
|
|
Size of the tensor.
|
|
Can be a tuple of one or two integers.
|
|
If one integer is given, the tensor will be square.
|
|
If two integers are given, the tensor will have the given width and height.
|
|
:param curvy: Float
|
|
Curviness of the gradient.
|
|
Higher values result in a sharper transition from 0 to 1.
|
|
:param scale: Float
|
|
Scale of the gradient.
|
|
Higher values result in a larger area with values close to 1.
|
|
:param mode: Str
|
|
Shape of the gradient.
|
|
Can be "square", "circle", "rectangle", or "corners".
|
|
:param min_val: Float
|
|
Minimum value in the tensor.
|
|
:param max_val: Float
|
|
Maximum value in the tensor.
|
|
:param center: Tuple of float
|
|
Coordinates of the center of the gradient.
|
|
:param function: Callable or None
|
|
Custom function for computing the distance from the center.
|
|
If given, this function will be used instead of the built-in modes.
|
|
The function should take two arguments (xx and yy) and return a tensor of the same shape.
|
|
:param normalize: Bool
|
|
Whether to normalize the tensor values to be between min_val and max_val.
|
|
:param dtype: Torch.dtype
|
|
Data type of the resulting tensor.
|
|
|
|
:return: Torch.Tensor
|
|
A 2D tensor with a radial gradient from 0 to 1 from the center to the edges.
|
|
"""
|
|
if isinstance(size, tuple):
|
|
if len(size) == 1:
|
|
width = height = size[0]
|
|
elif len(size) == 2:
|
|
width, height = size
|
|
else:
|
|
raise ValueError("Invalid size argument")
|
|
else:
|
|
raise TypeError("Size must be a tuple")
|
|
|
|
x = torch.linspace(0, 1, width)
|
|
y = torch.linspace(0, 1, height)
|
|
|
|
xx, yy = torch.meshgrid(x, y, indexing='ij')
|
|
|
|
if function is not None:
|
|
d = function(xx, yy)
|
|
elif mode == "square":
|
|
xx = (torch.abs(xx - center[0]) ** curvy)
|
|
yy = (torch.abs(yy - center[1]) ** curvy)
|
|
d = torch.max(xx, yy)
|
|
elif mode == "circle":
|
|
d = torch.sqrt((xx - center[0]) ** 2 + (yy - center[1]) ** 2)
|
|
d = (d ** curvy)
|
|
elif mode == "rectangle":
|
|
xx = (torch.abs(xx - center[0]) ** curvy)
|
|
yy = (torch.abs(yy - center[1]) ** curvy)
|
|
d = xx + yy
|
|
elif mode == "corners":
|
|
xx = (torch.abs(xx - center[0]) ** curvy)
|
|
yy = (torch.abs(yy - center[1]) ** curvy)
|
|
d = torch.min(xx, yy)
|
|
else:
|
|
raise ValueError("Not supported mode.")
|
|
|
|
if normalize:
|
|
d = d / d.max() * scale
|
|
d = torch.clamp(d, min_val, max_val)
|
|
|
|
return d.to(dtype)
|
|
|
|
|
|
Tensor.tensor_to_image = tensor_to_image
|
|
ImageB.image_to_tensor = image_to_tensor
|