Pytorch实现图像语义分割（初体验）

这些天在学习图像语义分割相关的知识，并简单写了篇概述。原本想先看几篇经典论文，如全卷积网络FCN，奈何英语水平有限，翻译起来实在费劲。想来不如先直接体验一下语义分割的效果，果然实践起来还挺有趣的。遂将过程记录如下。

代码实现

from torchvision import models
from PIL import Image
import matplotlib.pyplot as plt
import torch
import torchvision.transforms as T
import numpy as np# Define the helper function
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 == lr[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 rgbdef segment(net, path):img = Image.open(path)plt.imshow(img)plt.axis('off')plt.show()# Comment the Resize and CenterCrop for better inference resultstrf = 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)out = net(inp)['out']om = torch.argmax(out.squeeze(), dim=0).detach().cpu().numpy()rgb = decode_segmap(om)plt.imshow(rgb)plt.axis('off')plt.show()fcn = models.segmentation.fcn_resnet101(pretrained=True).eval()
# dlb = models.segmentation.deeplabv3_resnet101(pretrained=True).eval()girl = '../img/girl_dog.jpg'
segment(fcn, girl)
# segment(dlb, girl)