Deeplabv3-ResNet101由具有ResNet-101主干的Deeplabv3模型构成。预训练模型已在Pascal VOC数据集中存在的20个类别的COCO train2017子集中进行了训练。

这里利用Deeplabv3-ResNet101来进行语义分割获取人体蒙版，代码如下：

from torchvision import models
from PIL import Image
import matplotlib.pyplot as plt
import torch
import torchvision.transforms as T
import numpy as np
import os
import os.path as ospfile_path = '/root/Workspace/origin/vicon_03301_01'base_name = osp.basename(file_path)#vicon_03301_03color_dir = osp.join(file_path, 'Color')#color文件夹
if not osp.exists(color_dir):os.makedirs(color_dir)result_folder = '/root/Workspace/result'
output_folder = osp.join(result_folder, base_name)if not os.path.exists(output_folder):#就是在输出结果的文件夹下创建一个名为vicon_03301_03的文件夹os.makedirs(output_folder)def human_segment(net, path, nc=21):img = Image.open(path)trf = T.Compose([T.ToTensor(),#把一个取值范围是[0,255]的PIL.Image或者shape为(H,W,C)的numpy.ndarray，转换成形状为[C,H,W]，取值范围是[0,1.0]的torch.FloadTensorT.Normalize(mean=[0.485, 0.456, 0.406],#把tensor正则化，Normalized_image=(image-mean)/stdstd=[0.229, 0.224, 0.225])])inp = trf(img).unsqueeze(0)#返回一个新的张量，对输入的制定位置插入维度 1out = net(inp)['out']image = torch.argmax(out.squeeze(), dim=0).detach().cpu().numpy()label_colors = np.array([(255, 255, 255),  # 0=background# 1=aeroplane, 2=bicycle, 3=bird, 4=boat, 5=bottle(255, 255, 255), (255, 255, 255), (255, 255, 255), (255, 255, 255), (255, 255, 255),# 6=bus, 7=car, 8=cat, 9=chair, 10=cow(255, 255, 255), (255, 255, 255), (255, 255, 255), (255, 255, 255), (255, 255, 255),# 11=dining table, 12=dog, 13=horse, 14=motorbike, 15=person(255, 255, 255), (255, 255, 255), (255, 255, 255), (255, 255, 255), (0, 0, 0),# 16=potted plant, 17=sheep, 18=sofa, 19=train, 20=tv/monitor(255, 255, 255), (255, 255, 255), (255, 255, 255), (255, 255, 255), (255, 255, 255)])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)#堆栈save_image = osp.join(output_folder, osp.basename(path))plt.imsave(save_image, rgb)dlab = models.segmentation.deeplabv3_resnet101(pretrained=1).eval()for filename in os.listdir(color_dir):  #包含想要划分的图像的文件夹image_dir = osp.join(color_dir, filename)human_segment(dlab, image_dir)

实验效果

输入图片：

输出图片：