常常在工作之中遇到将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)