19 lines
588 B
Python
19 lines
588 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
|