29 lines
854 B
Python
29 lines
854 B
Python
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
|
|
def convert(image):
|
|
if isinstance(image, Image.Image):
|
|
img_array = np.array(image)
|
|
img_tensor = torch.from_numpy(img_array).float() / 255.
|
|
if img_tensor.ndim == 3 and img_tensor.shape[-1] == 3:
|
|
img_tensor = img_tensor.permute(2, 0, 1)
|
|
image = img_tensor.unsqueeze(0).permute(0, 2, 3, 1)
|
|
|
|
elif isinstance(image, torch.Tensor):
|
|
img_array = image.squeeze(0).cpu().numpy() * 255.0
|
|
image = Image.fromarray(np.clip(img_array, 0, 255).astype(np.uint8))
|
|
|
|
return image
|
|
|
|
|
|
# Tensor to PIL
|
|
def tensor2pil(image):
|
|
return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
|
|
|
|
|
|
# PIL to Tensor
|
|
def pil2tensor(image):
|
|
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
|