Numpy or Tensor array to Image

필요할때 마다 그때그때 코드 작성하기 너무 귀찮아서 정리했다.
Numpy array(H,W,C) 또는 Tensor array(C,H,W)에 대해서
알아서 numpy.ndarray로 변환하고 0~255로 normalized한 다음
input array에 대해 PIL.Image.Image를 반환하고 저장해주는 함수이다.
save_dir는 custom directory를 사용하면 된다.

def arr2imgsave(x, save_dir):
    # save the input image x and return PIL.Image.Image
    
    # Environment needed
    # from PIL import Image
    # import numpy as np
    # import torch
    
    # transform to numpy array
    if type(x) is torch.Tensor:
        x = x.permute(1,2,0).numpy()
    elif type(x) is np.ndarray:
        x = x
    else:
        print('Data type is not a numpy.ndarray or torch.tensor')
        return 0
    
    # numpy array normalize
    x = (x-x.min()) / (x.max() - x.min())
    x = (x * 255).astype(np.uint8)
    img = Image.fromarray(x)
    img.save(save_dir)
    
    return img