Convert between PIL Image and Numpy while using torchvision

image_converter.py

import numpy as np
from PIL import Image

__all__ = (
    'image_to_numpy',
    'numpy_to_image',
)

IMAGE_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32).reshape((3, 1, 1))
IMAGE_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32).reshape((3, 1, 1))


def image_to_numpy(img: Image) -> np.ndarray:
    return (np.array(img).transpose((2, 0, 1)).astype(np.float32) / 255 - IMAGE_MEAN) / IMAGE_STD


def numpy_to_image(arr: np.ndarray) -> Image:
    return Image.fromarray(
        ((arr * IMAGE_STD + IMAGE_MEAN) * 255).round().clip(0, 255).astype(np.uint8).transpose((1, 2, 0)))


if __name__ == '__main__':
    numpy_to_image(image_to_numpy(Image.open('./data/content.jpg'))).show()