将tensor转换为图像_tensor转int

import tensor 
import numpy as np
def tensor2img(tensor, out_type=np.uint8, min_max=(0, 1)):
''' Converts a torch Tensor into an image Numpy array Input: 4D(B,(3/1),H,W), 3D(C,H,W), or 2D(H,W), any range, RGB channel order Output: 3D(H,W,C) or 2D(H,W), [0,255], np.uint8 (default) '''
if hasattr(tensor, 'detach'):
tensor = tensor.detach()
tensor = tensor.squeeze().float().cpu().clamp_(*min_max)  # clamp
tensor = (tensor - min_max[0]) / (min_max[1] - min_max[0])  # to range [0,1]
n_dim = tensor.dim()
if n_dim == 4:
n_img = len(tensor)
img_np = make_grid(tensor, nrow=int(math.sqrt(n_img)), normalize=False).numpy()
img_np = np.transpose(img_np[[2, 1, 0], :, :], (1, 2, 0))  # HWC, BGR
elif n_dim == 3:
img_np = tensor.numpy()
img_np = np.transpose(img_np[[2, 1, 0], :, :], (1, 2, 0))  # HWC, BGR
elif n_dim == 2:
img_np = tensor.numpy()
else:
raise TypeError(
'Only support 4D, 3D and 2D tensor. But received with dimension: {:d}'.format(n_dim))
if out_type == np.uint8:
img_np = np.clip((img_np * 255.0).round(), 0, 255)
# Important. Unlike matlab, numpy.unit8() WILL NOT round by default.
return img_np.astype(out_type)