Torch的FCN-ResNet101语义分割模型是在COCO 2017训练集上的一个子集训练得到的，相当于PASCAL VOC数据集，支持20个类别。
from torchvision import models
from PIL import Image
import matplotlib.pyplot as plt
import torch
import torchvision.transforms as T
import numpy as np#获取模型，如果本地缓存没有，则会自动下载
fcn = models.segmentation.fcn_resnet101(pretrained=True).eval()#要预测的图像
img = Image.open('./bird.png')
plt.imshow(img)
plt.show()
在将图像数据输入网络之前，需要对图像进行预处理。
1. 将图像resize到256 x 256
2. 中心裁剪成224 x 224
3. 转换成Tensor归一化到[0,1]
4. 使用均值、方差标准化
5. Torch的输入数据格式是NCHW，所以还需要进行维度扩展
trf = T.Compose([T.Resize(256),
T.CenterCrop(224),
T.ToTensor(),
T.Normalize(mean = [0.485, 0.456, 0.406],
std = [0.229, 0.224, 0.225])])
inp = trf(img).unsqueeze(0)Torch模型预测的输出是一个OrderedDict结构。对于FCN-ResNet101这里使用的模型，它的输出大小是[1 x 21 x H x W]，21代表是20+1（background）个类别。
out = fcn(inp)['out']然后使用argmax选出每个类别概率最大的，并将第0个维度去掉，变成H x W的2D图像，这样每一个像素代表的就是该点的类别。
om = torch.argmax(out.squeeze(), dim=0).detach().cpu().numpy()最后将上面的输出映射成分割着色图。
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
rgb = decode_segmap(om)
plt.imshow(rgb)
plt.show()
