Spaces:
Runtime error
Runtime error
from PIL import Image | |
from torchvision import transforms | |
from skimage import io, transform, util | |
import numpy as np | |
import os | |
""" | |
Contains utility functions to work with images in tensor and jpg/png forms | |
""" | |
def load_image_tensor(image, path=""): | |
""" | |
Returns Image as a Pytorch Tensor of shape ((img_size),3). | |
Values between 0 and 1. | |
""" | |
img_size = (256, 256) | |
# image = io.imread(path) | |
cropped_image = util.crop(image, ((0, 0), (0, image.shape[1] - image.shape[0]), (0, 0))) | |
resized_image = (transform.resize(image=cropped_image, output_shape=img_size, anti_aliasing=True)) | |
to_tensor = transforms.Compose([transforms.ToTensor()]) | |
tensor = to_tensor(resized_image) | |
# tensor = tensor.permute(1,2,0) # the model expects w, h, 3! | |
return tensor.float() | |
def convert_tensor_to_PIL_image(image_tensor): | |
output_image = image_tensor.numpy().transpose(1, 2, 0) | |
output_image = np.clip(output_image, 0, 1) * 255 | |
output_image = output_image.astype(np.uint8) | |
output_image = Image.fromarray(output_image) | |
return output_image | |
def save_image_tensor(tensor, output_dir="./", image_name="output.png"): | |
""" | |
Saves a 3D tensor as an image. | |
""" | |
output_image = tensor.numpy().transpose(1, 2, 0) | |
output_image = np.clip(output_image, 0, 1) * 255 | |
output_image = output_image.astype(np.uint8) | |
output_image = Image.fromarray(output_image) | |
if not os.path.exists(output_dir): | |
os.mkdir(output_dir) | |
output_image.save(output_dir + image_name) | |
return output_image | |
def display_image_tensor(tensor): | |
""" | |
Displays the passed in 3D image tensor | |
""" | |
output_image = tensor.numpy().transpose(1, 2, 0) | |
output_image = np.clip(output_image, 0, 1) * 255 | |
output_image = output_image.astype(np.uint8) | |
output_image = Image.fromarray(output_image) | |
output_image.show() | |
def get_grayscale(tensor): | |
""" | |
Converts a 3D image tensor to greyscale | |
""" | |
greyscale_transform = transforms.Grayscale() | |
return greyscale_transform(tensor) | |