PyTorch基于预训练模型的图像分割预测示例

【图书推荐】《PyTorch深度学习与计算机视觉实践》-CSDN博客

PyTorch图像分割的基本形式-CSDN博客

下面将使用一个预训练的模型完成图像分割的示例。Torch的FCN-ResNet101语义分割模型是在COCO 2017训练集的一个子集上训练得到的，相当于PASCAL VOC数据集，支持20个类别。

调用FCN的代码如下：

from torchvision import models
获取模型，如果本地缓存没有，则会自动下载
fcn = models.segmentation.fcn_resnet101(pretrained=True).eval()

下面就是图像的准备，这里准备了一幅来自VOC动物数据集的图像，如图14-5所示。

图14-5  准备的图像

有了图像，自然而然的想法就是使用FCN对这幅图进行图像分割。但是在将它输入模型之前，还需要进行预处理，顺序如下：

（1）将图像大小调整到224×224。

（2）转换成Tensor，归一化到[0,1]。

（3）使用均值、方差标准化。

（4）Torch的输入数据格式是NCHW，所以还需要进行维度扩展。

这部分处理的代码如下：

import torchvision.transforms as T
transform = T.Compose([T.ToTensor(),T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
from_img = cv2.imread("./from_img.jpg");        
from_img = cv2.cvtColor(from_img,cv2.COLOR_BGR2RGB)
from_img = cv2.resize(from_img, (224, 224))
from_img = transform(from_img).unsqueeze(0)

Torch模型预测的输出是一个OrderedDict结构。对于FCN-ResNet101，它的输出大小是[1×21×H×W]，21代表是20+1（background）个类别。然后使用argmax选出每个类别中概率最大的，并将第0个维度去掉，变成H×W的二维图像，这样每一个像素代表的就是该点的类别。

out = fcn(from_img)['out']
print (out.shape)

import numpy as np
om = torch.argmax(out.squeeze(), dim=0).detach().cpu().numpy()
print (om.shape)

print (np.unique(om))

打印结果如下：

可以看到，通过FCN模型计算后，数据根据类别计算成[1,31,224,224]大小的矩阵向量，之后经过处理生成一个和原图像大小相同的二维矩阵。

out是模型的最终输出，其形状是[1×21×H×W]，与前面所讨论的结果一致。因为模型是在21个类（包括背景类）上训练的，输出有21个通道，其中该图像的每个像素对应于一个类，所以二维图像（形状[H×W]）的每个像素将与相应的类标签对应。对于该二维图像中的每个（x，y）像素，将对应于表示类的0~20的数字。

Unique函数的作用就是将矩阵的数值按“类”提取出来，从打印的结果来看，处理后的图像只包含两种元素，即0和17。现在有了一个二维图像，其中每个像素属于一个类。

下面一步就是将预测的图像解码输出，可以使用如下代码将二维图像转换为RGB图像，其中每个（元素）标签映射到相应的颜色。

def decode_segmap(image, nc=21):
    label_colors = np.array([(0, 0, 0),  #0=background
                             #1=aeroplane, 2=bicycle, 3=bird, 4=boat, 5=bottle
                             (128, 0, 0), (0, 128, 0), (128, 128, 0), (0, 0, 128), (128, 0, 128),
                             #6=bus, 7=car, 8=cat, 9=chair, 10=cow
                             (0, 128, 128), (128, 128, 128), (64, 0, 0), (192, 0, 0), (64, 128, 0),
                             #11=dining table, 12=dog, 13=horse, 14=motorbike, 15=person
                             (192, 128, 0), (64, 0, 128), (192, 0, 128), (64, 128, 128), (192, 128, 128),
                             #16=potted plant, 17=sheep, 18=sofa, 19=train, 20=tv/monitor
                             (0, 64, 0), (128, 64, 0), (0, 192, 0), (128, 192, 0), (0, 64, 128)])

    r = np.zeros_like(image).astype(np.uint8)
    g = np.zeros_like(image).astype(np.uint8)
    b = np.zeros_like(image).astype(np.uint8)

    for l in range(0, nc):
        idx = image == l
        r[idx] = label_colors[l, 0]
        g[idx] = label_colors[l, 1]
        b[idx] = label_colors[l, 2]

    rgb = np.stack([r, g, b], axis=2)
    return rgb

最后将图像用openCV进行展示，代码如下：

rgb = decode_segmap(om)
cv2.imshow("seg_img",rgb)
cv2.waitKey(0)

最终结果如图14-6所示。

图14-6  输出结果

这是经过模型预测后的模型输出，可以看到，对于原始图像，已经可以较为明显地分辨出图像中存在的生物，虽然在细节上还是欠佳，但是作为图像分割的基准还是可以使用的。