38 lines
1006 B
Python
38 lines
1006 B
Python
import io
|
|
import torch
|
|
import base64
|
|
import numpy as np
|
|
from pkg_resources import parse_version
|
|
from PIL import Image
|
|
|
|
|
|
def pil2numpy(image: Image.Image):
|
|
return np.array(image).astype(np.float32) / 255.0
|
|
|
|
|
|
def numpy2pil(image: np.ndarray, mode=None):
|
|
return Image.fromarray(np.clip(255.0 * image, 0, 255).astype(np.uint8), mode)
|
|
|
|
|
|
## Helper function equivalent to Mikey's pil2tensor
|
|
#def pil2tensor(self, image):
|
|
# return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
|
|
|
|
def pil2tensor(image: Image.Image):
|
|
return torch.from_numpy(pil2numpy(image)).unsqueeze(0)
|
|
|
|
|
|
def tensor2pil(image: torch.Tensor, mode=None):
|
|
return numpy2pil(image.cpu().numpy().squeeze(), mode=mode)
|
|
|
|
|
|
def tensor2bytes(image: torch.Tensor) -> bytes:
|
|
return tensor2pil(image).tobytes()
|
|
|
|
|
|
def pil2base64(image: Image.Image):
|
|
buffered = io.BytesIO()
|
|
image.save(buffered, format="PNG")
|
|
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
|
|
return img_str
|