将torch.Tensor格式的图片转换为numpy.ndarray格式

下述两个函数均可:
1.

def image_Tensor2ndarray(image_tensor: torch.Tensor):
    """
    将tensor转化为cv2格式
    """
    assert (len(image_tensor.shape) == 4 and image_tensor.shape[0] == 1)
    # 复制一份
    image_tensor = image_tensor.clone().detach()
    # 到cpu
    image_tensor = image_tensor.to(torch.device('cpu'))
    # 反归一化
    # input_tensor = unnormalize(input_tensor)
    # 去掉批次维度
    image_tensor = image_tensor.squeeze()
    # 从[0,1]转化为[0,255],再从CHW转为HWC,最后转为cv2
    image_cv2 = image_tensor.mul_(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).type(torch.uint8).numpy()
    # RGB转BRG
    # image_cv2 = cv2.cvtColor(image_cv2, cv2.COLOR_BGR2RGB)
    return image_cv2

def image_Tensor2np(image_tensor: torch.Tensor):
    """
    将tensor转化为cv2格式
    """
    assert (len(image_tensor.shape) == 4 and image_tensor.shape[0] == 1)

    # 复制一份
    image_tensor = image_tensor.cpu().detach()
    # 去掉批次维度
    image_tensor = image_tensor.squeeze()
    # 将tensor数据转为numpy数据
    image_tensor = image_tensor.numpy()
    # 反归一化normalize,将图像数据扩展到[0,255]
    maxValue = image_tensor.max()
    image_tensor = image_tensor * 255 / maxValue
    # float32-->uint8
    image_cv2 = np.uint8(image_tensor)
    # 再从CHW转为HWC
    image_cv2 = image_cv2.transpose(1, 2, 0)
    # # RGB转BRG
    # image_cv2 = cv2.cvtColor(image_cv2, cv2.COLOR_RGB2BGR)
    return image_cv2

你可能感兴趣的:(python,pytorch,python)