常常在工作之中遇到將dataloader中出來的tensor成image,numpy格式的數據,然后可以可視化出來
但是這種tensor往往經過了channel變換(RGB2BGR),以及歸一化(減均值除方差),
然后維度的順序也發生變化(HWC變成CHW)。為了可視化這種變化比較多的數據,
在tensor轉numpy之前需要對tensor做一些處理
如下是一個簡單的函數,可以可視化tensor,下次直接拿來用就行
def tensor2im(input_image, imtype=np.uint8):
""""
Parameters:
input_image (tensor) -- 輸入的tensor,維度為CHW,注意這里沒有batch size的維度
imtype (type) -- 轉換后的numpy的數據類型
"""
mean = [0.485, 0.456, 0.406] # dataLoader中設置的mean參數,需要從dataloader中拷貝過來
std = [0.229, 0.224, 0.225] # dataLoader中設置的std參數,需要從dataloader中拷貝過來
if not isinstance(input_image, np.ndarray):
if isinstance(input_image, torch.Tensor): # 如果傳入的圖片類型為torch.Tensor,則讀取其數據進行下面的處理
image_tensor = input_image.data
else:
return input_image
image_numpy = image_tensor.cpu().float().numpy() # convert it into a numpy array
if image_numpy.shape[0] == 1: # grayscale to RGB
image_numpy = np.tile(image_numpy, (3, 1, 1))
for i in range(len(mean)): # 反標准化,乘以方差,加上均值
image_numpy[i] = image_numpy[i] * std[i] + mean[i]
image_numpy = image_numpy * 255 #反ToTensor(),從[0,1]轉為[0,255]
image_numpy = np.transpose(image_numpy, (1, 2, 0)) # 從(channels, height, width)變為(height, width, channels)
else: # 如果傳入的是numpy數組,則不做處理
image_numpy = input_image
return image_numpy.astype(imtype)