21 lines
502 B
Python
21 lines
502 B
Python
import torchvision.transforms as t
|
|
from PIL import Image as ImageF
|
|
from PIL.Image import Image as ImageB
|
|
from torch import Tensor
|
|
|
|
|
|
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)
|
|
|
|
|
|
Tensor.tensor_to_image = tensor_to_image
|
|
ImageB.image_to_tensor = image_to_tensor
|