OpenCV图像到PyTorch张量的预处理

OpenCV图像读取后，默认是预处理为HWC形状，而在PyTorch中，希望的是CHW形状; 另一方面，在通道顺序上，OpenCV默认是BGR的顺序，而在PyTorch的张量处理环节，我们一般希望的是RGB顺序。这里就涉及到了数据处理和转换。

从OpenCV图像到PyTorch张量的转换

我们先模拟一个高度为2，宽度为3的彩色图像：

from torch import tensor

a = torch.randint(low=0, high=255, size=(2, 3, 3))

这里的a似于：

tensor([[[ 58, 223,  24], [ 35, 116, 249], [ 81, 220, 168]],
        [[242, 169, 149], [149,  87, 167], [165,  66,  88]]])

如果这个数据是从OpenCV读取过来，那么最内层通道顺序是BGR。

为了得到RGB，我们可以将最内侧倒排一下：

b = a[:, :, [2,1,0]]

结果类似于：

tensor([[[ 24, 223,  58], [249, 116,  35], [168, 220,  81]],
        [[149, 169, 242], [167,  87, 149], [ 88,  66, 165]]])

但问题是，上面的张量形状是 (H,W,C)，而在PyTorch中，我们期望的是(C,H,W)。现在我们将它拆分重排成RGB三个通道，其中每个通道都是一个2*3灰度图：

c = b.permute(2,0,1 )
print(c)

结果类似于：

tensor([
        # R通道
        [[ 24, 249, 168],
         [149, 167,  88]],         
        # G通道
        [[223, 116, 220],
         [169,  87,  66]],
        # B通道
        [[ 58,  35,  81],
         [242, 149, 165]]])

一个数据集加载示例

在《Modern Computer Vision》中，讲解了使用VGG16进行迁移学习，然后用于猫狗分类的示例。一方面VGG16预训练模型是在ImageNet数据集上训练得来的，其中所有图像都被缩放为 224 × 224 作为输入。为了充分利用预训练的权重，我们需要把训练的输入图像调整为相同大小。

另一方面，由于VGG期望的归一化参数是：

• mean: [0.485, 0.456, 0.406]
• std: [0.229, 0.224, 0.225]

所以，我们的训练图像也需要按上面的参数进行归一化。下面的代码完成了上述的转换：

from torch import tensor


class CatsDogs(Dataset):

    def __init__(self, folder):
        cats = glob(folder+'/cats/*.jpg')
        dogs = glob(folder+'/dogs/*.jpg')
        self.fpaths = cats[:500] + dogs[:500]
        self.normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])
        from random import shuffle, seed; 
        seed(10); 
        shuffle(self.fpaths)
        self.targets = [
            fpath.split('/')[-1].startswith('dog') 
            for fpath in self.fpaths
            ] 

    def __len__(self): 
        return len(self.fpaths)
    
    def __getitem__(self, ix):
        f = self.fpaths[ix]
        target = self.targets[ix]
        im = cv2.imread(f)
        im = cv2.resize(im, (224,224))    # 尺寸调整
        im = torch.tensor(im/255.0)       # 张量化
        im = im[:,:,[2,1,0]]              # RGB -> BGR
        im = im.permute(2,0,1)            # HWC -> CHW
        im = self.normalize(im)           # 归一化

        x = im.float().to(device) 
        y = tensor([target]).float().to(device)
        return x, y

以上代码是在Modern-Computer-Vision-with-PyTorch的基础上改动而来。主要是修改了通道顺序调整方式（从numpy上改动放到了在tensor上改动）和位置（和通道顺序调整放到了一起）。