PyTorch: 各种图像格式相互转化

淡淡的烟草味﹌ 2022-11-12 08:30 252阅读 0赞

# PyTorch: 各种图像格式相互转化 #

我们一般在PyTorch或者Python中处理的图像无非这几种格式：

*  PIL：使用Python自带图像处理库读取出来的图片格式;
 *  numpy：使用python-opencv库读取出来的图片格式;
 *  tensor：pytorch中训练时所采取的向量格式（当然也可以说图片）.

注意，本文所述图片格式皆为RGB三通道，24-bit真彩色，也就是我们平常使用的图片形式。

下述代码均引用了：

import cv2
    import torch
    from PIL import Image
    import matplotlib.pyplot as plt
    
    # loader使用torchvision中自带的transforms函数
    loader = transforms.Compose([
        transforms.ToTensor()])  
    
    unloader = transforms.ToPILImage()

## PIL读取图片转化为Tensor ##

# 输入图片地址, 返回tensor变量
    def image_loader(image_name):
        image = Image.open(image_name).convert('RGB')
        image = loader(image).unsqueeze(0)
        return image.to(device, torch.float)

# 输入PIL格式图片, 返回tensor变量
    def PIL_to_tensor(image):
        image = loader(image).unsqueeze(0)
        return image.to(device, torch.float)

## Tensor转化为PIL图片 ##

# 输入tensor变量, 输出PIL格式图片
    def tensor_to_PIL(tensor):
        image = tensor.cpu().clone()
        image = image.squeeze(0)
        image = unloader(image)
        return image

## 直接展示Tensor格式图片 ##

def imshow(tensor, title=None):
        image = tensor.cpu().clone()  # we clone the tensor to not do changes on it
        image = image.squeeze(0)  # remove the fake batch dimension
        image = unloader(image)
        plt.imshow(image)
        if title is not None:
            plt.title(title)
        plt.pause(0.001)  # pause a bit so that plots are updated

## 直接保存Tensor格式图片 ##

def save_image(tensor, **para):
        dir = 'results'
        image = tensor.cpu().clone()  # we clone the tensor to not do changes on it
        image = image.squeeze(0)  # remove the fake batch dimension
        image = unloader(image)
        if not osp.exists(dir):
            os.makedirs(dir)
        image.save('results_{}/s{}-c{}-l{}-e{}-sl{:4f}-cl{:4f}.jpg'
                   .format(num, para['style_weight'], para['content_weight'], para['lr'], para['epoch'],
                           para['style_loss'], para['content_loss']))

## Numpy转化为Tensor ##

def toTensor(img):
        assert type(img) == np.ndarray,'the img type is {}, but ndarry expected'.format(type(img))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = torch.from_numpy(img.transpose((2, 0, 1)))
        return img.float().div(255).unsqueeze(0)  # 255也可以改为256

## Tensor转化为Numpy ##

def tensor_to_np(tensor):
        img = tensor.mul(255).byte()
        img = img.cpu().numpy().squeeze(0).transpose((1, 2, 0))
        return img

## 展示Numpy格式图片 ##

def show_from_cv(img, title=None):
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        plt.figure()
        plt.imshow(img)
        if title is not None:
            plt.title(title)
        plt.pause(0.001)

## 展示Tensor格式图片 ##

def show_from_tensor(tensor, title=None):
        img = tensor.clone()
        img = tensor_to_np(img)
        plt.figure()
        plt.imshow(img)
        if title is not None:
            plt.title(title)
        plt.pause(0.001)

注意: 上面介绍的都是一张图片的转化，如果是n张图片一起的话，只需要修改一下相应代码即可。

# 将 N x H x W X C 的numpy格式图片转化为相应的tensor格式
    def toTensor(img):
        img = torch.from_numpy(img.transpose((0, 3, 1, 2)))
        return img.float().div(255).unsqueeze(0)