针对图像的语义分割网络,介绍PyTorch中已经训练好的网络的使用方式。
在PyTorch提供的已训练好的图像语义分割网络中,有两类预训练好的网络,分别是FCNResNet101系列和DeepLapV3 ResNet101系列。针对语义分割的分类器,需要输入图像使用了相同的预处理方式,即先将每张图像的像素值预处理到0~1之间,然后对图像进行标准化处理。
以segmentation.fcn_resnet101()为例,具体代码如下:
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import PIL
import torch
from torchvision import transforms
import torchvision
model = torchvision.models.segmentation.fcn_resnet101(pretrained=True)
model.eval()
image = PIL.Image.open("F:/test/test_photo/23.jpg") #图片路径##
image_transf = transforms.Compose([
transforms.ToTensor(),transforms.Normalize(mean = [0.485,0.456,0.406],
std = [0.229,0.224,0.225])
])
image_tensor = image_transf(image).unsqueeze(0)
output= model(image_tensor)["out"]
outputarg = torch.argmax(output.squeeze(),dim=0).numpy()
outputarg
def decode_segmaps(image,Label_colors,nc=21):
"""函数将会输出2D图像,会将不同的内编码为不同的颜色"""
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 cla in range(0,nc):
idx = image ==cla
r[idx] = Label_colors[cla,0]
g[idx] = Label_colors[cla,1]
b[idx] = Label_colors[cla,2]
rgbimage = np.stack([r,g,b],axis=2)
return rgbimage
label_colors = np.array([(0,0,0),
(128,0,0),(0,128,0),(128,128,0),(0,0,128),(128,0,128),
(0,128,128),(128,128,128),(64,0,0),(192,0,0),(64,128,0),
(129,128,0),(64,0,128),(192,0,128),(64,128,128),(192,128,128),
(0,64,0),(128,64,0),(0,192,0),(128,192,0),(0,64,128)])
outputrgb =decode_segmaps(outputarg,label_colors)
plt.figure(figsize=(20,8))
plt.subplot(1,2,1)
plt.imshow(image)
plt.axis("off")
plt.subplot(1,2,2)
plt.imshow(outputrgb)
plt.axis("off")
plt.subplots_adjust(wspace=0.05)
plt.show()
经测试,发现预训练的网络对于常见人物、动物分割效果比较理想!