Pytorch 语义分割DeepLabV3+ 训练自己的数据集

照葫芦画瓢总结记录了一下DeepLab分割系列，并对Deeplab V3++实现

一、DeepLab系列理解

1、DeepLab V1

原文：Semantic image segmentation with deep convolutional nets and fully connected CRFs（https://arxiv.org/pdf/1412.7062v3.pdf）
收录：ICLR 2015 (International Conference on Learning Representations)
Backbone： VGG16
Contributions：
Atrous convolution
CRF
DeepLab V1是基于VGG16网络改写的，一共做了三件事。

首先，去掉了最后的全连接层。做语义分割使用全卷积网络是大势所趋。

然后，去掉了最后两个池化层。池化层是神经网络中的一个经典结构，BP解决了神经网络训练的软件问题（权重更新），pooling解决了训练的硬件问题（对计算资源的需求）。这就是池化层的第一个作用，缩小特征层的尺寸。池化层还有另一个重要作用，快速扩大感受野。为什么要扩大感受野呢？为了利用更多的上下文信息进行分析。

在实验中发现 DCNNs 做语义分割时精准度不够的问题，根本原因是 DCNNs 的高级特征的平移不变性，即高层次特征映射，根源于重复的池化和下采样。
针对信号下采样或池化降低分辨率，DeepLab 采用的空洞卷积算法扩展感受野，获取更多的语境信息。

语义分割是一个end-to-end的问题，需要对每个像素进行精确的分类，对像素的位置很敏感，pooling是一个不断丢失位置信息的过程，而语义分割又需要这些信息，矛盾就产生了。没办法，只好去掉pooling喽。全去掉行不行，理论上是可行的，实际使用嘛，一来显卡没那么大的内存，二来费时间。所以只去掉了两层。
（PS：在DeepLab V1原文中，作者还指出一个问题，使用太多的pooling，特征层尺寸太小，包含的特征太稀疏了，不利于语义分割）

去了两个pooling，感受野又不够了怎么办？把atrous convolution借来用一下，这也是对VGG16的最后一个修改。atrous convolution人称空洞卷积（好像多称为dilation convolution），相比于传统卷积，可以在不增加计算量的情况下扩大感受野。

空洞卷积与传统卷积的区别在于，传统卷积是三连抽，感受野是3，空洞卷积是跳着抽，也就是使用图中的rate，感受野一下扩大到了5（rate=2），相当于两个传统卷积，而通过调整rate可以自由选择感受野。这样感受野的问题就解决了。

另外，原文指出，空洞卷积的优势在于增加了特征的密度。这张图你不能单独看，上边的传统卷积是经过pooling以后的第一个卷积层，而下边卷积输入的浅粉色三角正是被pooling掉的像素。所以，下边的输出是上边的两倍，特征多出了一倍。

DeepLab V1的另一个贡献是使用条件随机场CRF提高分类精度。效果如下图，可以看到提升是非常明显的。具体CRF是什么原理呢？没有去研究，因为懒到了V3就舍弃了CRF。采用完全连接的条件随机场（CRF）提高模型捕获细节的能力，简单来说，就是对一个像素进行分类的时候，不仅考虑DCNN的输出，而且考虑该像素点周围像素点的值，这样语义分割结果边界清楚。

除空洞卷积和 CRFs 之外，论文使用的 tricks 还有 Multi-Scale features，与FCN skip layer类似，具体实现上，在输入图片与前四个 max pooling 后添加卷积层，这四个预测结果和模型输出拼接。

DeepLab V1的全流程图，其中上采样直接使用了双线性采样。

2、DeepLab V2

原文：DeepLab: Semantic Image Segmentation with Deep Convolutional Nets, Atrous Convolution, and Fully Connected CRFs（ https://arxiv.org/abs/1606.00915）
收录：TPAMI2017 (IEEE Transactions on Pattern Analysis and Machine Intelligence, 2017)
Backbone：ResNet-101
Contributions：
ASPP
在DeepLab V2中，可能是觉得VGG16表达能力有限，于是大神换用了更复杂，表达能力更强的ResNet-101

在V2中，大神同样对ResNet动了刀，刀法和V1相同。V2的贡献在于更加灵活的使用了atrous convolution，提出了空洞空间金字塔池化ASPP。还是先给图

ASPP的作用，说白了就是利用空洞卷积的优势，从不同的尺度上提取特征。这么做的原因也很简单，因为相同的事物在同一张图或不同图像中存在尺度上的差异。还是以这张图为例，图中的树存在多种尺寸，使用ASPP就能更好的对这些树进行分类。

至于ASPP如何融合到ResNet中，将VGG16的conv6，换成不同rate的空洞卷积，再跟上conv7，8，最后做个大融合（对应相加或1*1卷积）就OK了。

3、DeepLab V3

原文：Rethinking Atrous Convolution for Semantic Image Segmentation（ https://arxiv.org/abs/1706.05587）
Backbone：ResNet-101
Contributions：
ASPP
Going deeper with atrous convolution
Remove CRF

舍弃了CRF，因为分类结果精度已经提高到不需要CRF了。
另外两个贡献，一个是改进了ASPP，另一个是使用空洞卷积加深网络，这两者算是一个二选一吧，是拓展网络的宽度，还是增加网络的深度。一般说起DeepLab V3模型指的是前者，因为从大神给出的结果和后续发展来看，明显前者效果更好一些。

对于ASPP，V3中做了两点改进。一是在空洞卷积之后使用batch normalization，对训练很有帮助。第二点是增加了1∗11*11∗1卷积分支和image pooling分支。增加这两个分支是为了解决使用空洞卷积带来的问题，随着rate的增大，一次空洞卷积覆盖到的有效像素（特征层本身的像素，相应的补零像素为非有效像素）会逐渐减小到1。这就与我们的初衷（获取更大范围的特征）相背离了。所以为了解决这个问题，一是使用1∗11*11∗1的卷积，也就是当rate增大以后3∗33*33∗3卷积的退化形式，替代3∗33*33∗3卷积，减少参数个数；另一点就是增加image pooling，可以叫做全局池化，来补充全局特征。具体做法是对每一个通道的像素取平均，之后再上采样到原来的分辨率。

ASPP的变化就这么多，再来简单说说另一种思路Going deeper with atrous convolution。为什么要加深网络呢，我理解的是为了获取更大的感受野。提到感受野自然离不开空洞卷积。还是看图说话，很显然pooling用多了特征层都快小的看不见了，所以大神给出了使用空洞卷积不断加深网络的一种思路。

4、DeepLab V3+

原文：Encoder-Decoder with Atrous Separable Convolution for Semantic Image Segmentation（ https://arxiv.org/abs/1802.02611）
Backbone：Xception
Contributions：
Xception
Encoder-decoder structure

DeepLab V3+再次修改了主网络，将ResNet-101升级到了Xception。在原始的Xception的基础上，大神进行了三点修改：
（1）使用更深的网络；
（2）将所有的卷积层和池化层用深度分离卷积Depthwise separable convolution进行替代，也就是下图中的Sep Conv；
（3）在每一次3*3 depthwise convolution之后使用BN和ReLU。

因为V3+使用深度分离卷积替代了pooling，那么为了缩小特征层尺寸，有几个block的最后一层的stride就必须为2，也就是下图中标红的层。具体有几个取决于输出output stride（下采样的大小）的设置。

Xception的核心是使用了Depthwise separable convolution。Depthwise separable convolution的思想来自inception结构，是inception结构的一种极限情况。Inception 首先给出了一种假设：卷积层通道间的相关性和空间相关性是可以退耦合的，将它们分开映射，能达到更好的效果。在inception结构中，先对输入进行1∗11*11∗1的卷积，之后将通道分组，分别使用不同的3∗33*33∗3卷积提取特征，最后将各组结果串联在一起作为输出。

Depthwise separable convolution是将这种分组演化到了极致，即把每一个通道作为一组。先对输入的每一个通道做3∗33*33∗3的卷积，将各个通道的结果串联后，再通过1∗11*11∗1的卷积调整到目标通道数。

使用depthwise separable convolution的好处。好处也很简单，大幅缩减参数个数。假设输入输出都是64通道，卷积核采用3*3，那么传统卷积的参数个数为
3∗3∗64∗64=36864

而Depthwise separable convolution为

3 ∗ 3 ∗ 64 + 1 ∗ 1 ∗ 64 ∗ 64 = 4672 3364+1164*64=4672
3∗3∗64+1∗1∗64∗64=4672

说完了backbone，再来说说V3+的整体结构。前三个版本都是backbone（ASPP）输出的结果直接双线性上采样到原始分辨率，非常简单粗暴的方法，下图中的(a)。用了三个版本，大神也觉得这样做太粗糙了，于是吸取Encoder-Deconder的结构，下图中的(b)，增加了一个浅层到输出的skip层，下图中的( c)。

具体的skip方法。首先，选取block2中的第二个卷积输出（看代码这个是固定的），使用1∗11*11∗1卷积调整通道数到48（减小通道数是为了降低其在最终结果中的比重），然后resize到指定的尺寸，也就是output stride。然后，将ASPP的输出resize到output stride。最后将两部分串联起来做两次3∗33*33∗3的卷积。最后的最后再做一次1*1的卷积，得到分类结果。最后的最后的最后将分类结果resize到原来的分辨率，采用双线性采样。
完整的流程图：

二、DeepLabV3+实现（cuda9.0，pytorch1.5）

1、数据：
之前用labellme打过标签的Abyssinian数据集
如何制作分割数据集见：https://blog.csdn.net/yx868yx/article/details/105641947（标注、生成label.png）和
https://blog.csdn.net/yx868yx/article/details/105642324（所有json文件夹中提取所有的label.png图片）
数据集形式：

		ImageSets	Segmentation(train.txt,val.txt)
VOCdevkit	Abyssinian	JPEGImages	训练和验证的图片
		SegmentationClass	训练和验证的mask图片

如图：

2、环境配置：
环境 ubantu16.04+cudnn7.0+cuda_9.0.176
Pytorch 1.5.1
3、clone工程： https://github.com/jfzhang95/pytorch-deeplab-xception
4、修改mypath.py
添加上自己的数据集名字和所在的路径。代码如下

class Path(object):@staticmethoddef db_root_dir(dataset):if dataset == 'pascal':return '/path/to/datasets/VOCdevkit/VOC2012/' # folder that contains VOCdevkit/.elif dataset == 'sbd':return '/path/to/datasets/benchmark_RELEASE/' # folder that contains dataset/.elif dataset == 'cityscapes':return '/path/to/datasets/cityscapes/' # foler that contains leftImg8bit/elif dataset == 'coco':return '/path/to/datasets/coco/'#添加我自己制作的Abyssinian数据集elif dataset == 'Abyssinian':return '/home/yuxin/pytorch-deeplab-xception-master/VOCdevkit/Abyssinian'else:print('Dataset {} not available.'.format(dataset))raise NotImplementedError

5 、创建Abyssinian.py文件
然后在dataloaders/datasets路径下创建自己的数据集文件，这里创建为Abyssinian.py，因为我们是按照VOC数据集格式，所以直接将pascal.py内容复制过来修改，这里将类别数和数据集改成自己的数据集，类别记得加上背景类，我这里一共是2类：

from __future__ import print_function, divisionimport os
from PIL import Image
import numpy as np
from torch.utils.data import Dataset
from mypath import Path
from torchvision import transforms
from dataloaders import custom_transforms as tr
class VOCSegmentation(Dataset):"""PascalVoc dataset"""#Abyssinian类和背景类NUM_CLASSES = 2def __init__(self,args,#修改为Abyssinianbase_dir=Path.db_root_dir('Abyssinian'),split='train',):""":param base_dir: path to VOC dataset directory:param split: train/val:param transform: transform to apply"""super().__init__()self._base_dir = base_dirself._image_dir = os.path.join(self._base_dir, 'JPEGImages')self._cat_dir = os.path.join(self._base_dir, 'SegmentationClass')if isinstance(split, str):self.split = [split]else:split.sort()self.split = splitself.args = args_splits_dir = os.path.join(self._base_dir, 'ImageSets', 'Segmentation')self.im_ids = []self.images = []self.categories = []for splt in self.split:with open(os.path.join(os.path.join(_splits_dir, splt + '.txt')), "r") as f:lines = f.read().splitlines()for ii, line in enumerate(lines):_image = os.path.join(self._image_dir, line + ".jpg")_cat = os.path.join(self._cat_dir, line + ".png")assert os.path.isfile(_image)assert os.path.isfile(_cat)self.im_ids.append(line)self.images.append(_image)self.categories.append(_cat)assert (len(self.images) == len(self.categories))# Display statsprint('Number of images in {}: {:d}'.format(split, len(self.images)))def __len__(self):return len(self.images)def __getitem__(self, index):_img, _target = self._make_img_gt_point_pair(index)sample = {'image': _img, 'label': _target}for split in self.split:if split == "train":return self.transform_tr(sample)elif split == 'val':return self.transform_val(sample)def _make_img_gt_point_pair(self, index):_img = Image.open(self.images[index]).convert('RGB')_target = Image.open(self.categories[index])return _img, _targetdef transform_tr(self, sample):composed_transforms = transforms.Compose([tr.RandomHorizontalFlip(),tr.RandomScaleCrop(base_size=self.args.base_size, crop_size=self.args.crop_size),tr.RandomGaussianBlur(),tr.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),tr.ToTensor()])return composed_transforms(sample)def transform_val(self, sample):composed_transforms = transforms.Compose([tr.FixScaleCrop(crop_size=self.args.crop_size),tr.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),tr.ToTensor()])return composed_transforms(sample)def __str__(self):return 'VOC2012(split=' + str(self.split) + ')'if __name__ == '__main__':from dataloaders.utils import decode_segmapfrom torch.utils.data import DataLoaderimport matplotlib.pyplot as pltimport argparseparser = argparse.ArgumentParser()args = parser.parse_args()args.base_size = 513args.crop_size = 513voc_train = VOCSegmentation(args, split='train')dataloader = DataLoader(voc_train, batch_size=5, shuffle=True, num_workers=0)for ii, sample in enumerate(dataloader):for jj in range(sample["image"].size()[0]):img = sample['image'].numpy()gt = sample['label'].numpy()tmp = np.array(gt[jj]).astype(np.uint8)segmap = decode_segmap(tmp, dataset='pascal')img_tmp = np.transpose(img[jj], axes=[1, 2, 0])img_tmp *= (0.229, 0.224, 0.225)img_tmp += (0.485, 0.456, 0.406)img_tmp *= 255.0img_tmp = img_tmp.astype(np.uint8)plt.figure()plt.title('display')plt.subplot(211)plt.imshow(img_tmp)plt.subplot(212)plt.imshow(segmap)if ii == 1:breakplt.show(block=True)

6、修改dataloaders/utils.py
在文件中创建一个get_Abyssinian_labels()函数，根据自己的类别数，我这里只有两类，一类是背景，一类是我要分割的类，所以只定义了两种mask颜色

from __future__ import print_function, division
import os
from PIL import Image
import numpy as np
from torch.utils.data import Dataset
from mypath import Path
from torchvision import transforms
from dataloaders import custom_transforms as tr
import matplotlib.pyplot as plt
import numpy as np
import torchdef decode_seg_map_sequence(label_masks, dataset='pascal'):rgb_masks = []for label_mask in label_masks:rgb_mask = decode_segmap(label_mask, dataset)rgb_masks.append(rgb_mask)rgb_masks = torch.from_numpy(np.array(rgb_masks).transpose([0, 3, 1, 2]))return rgb_masksdef decode_segmap(label_mask, dataset, plot=False):"""Decode segmentation class labels into a color imageArgs:label_mask (np.ndarray): an (M,N) array of integer values denotingthe class label at each spatial location.plot (bool, optional): whether to show the resulting color imagein a figure.Returns:(np.ndarray, optional): the resulting decoded color image."""if dataset == 'pascal' or dataset == 'coco':n_classes = 21label_colours = get_pascal_labels()elif dataset == 'cityscapes':n_classes = 19label_colours = get_cityscapes_labels()#添加Abyssinian elif dataset ==‘Abyssinian':n_classes = 2label_colours = get_Abyssinian_labels()else:raise NotImplementedErrorr = label_mask.copy()g = label_mask.copy()b = label_mask.copy()for ll in range(0, n_classes):r[label_mask == ll] = label_colours[ll, 0]g[label_mask == ll] = label_colours[ll, 1]b[label_mask == ll] = label_colours[ll, 2]rgb = np.zeros((label_mask.shape[0], label_mask.shape[1], 3))rgb[:, :, 0] = r / 255.0rgb[:, :, 1] = g / 255.0rgb[:, :, 2] = b / 255.0if plot:plt.imshow(rgb)plt.show()else:return rgbdef encode_segmap(mask):"""Encode segmentation label images as pascal classesArgs:mask (np.ndarray): raw segmentation label image of dimension(M, N, 3), in which the Pascal classes are encoded as colours.Returns:(np.ndarray): class map with dimensions (M,N), where the value ata given location is the integer denoting the class index."""mask = mask.astype(int)label_mask = np.zeros((mask.shape[0], mask.shape[1]), dtype=np.int16)for ii, label in enumerate(get_pascal_labels()):label_mask[np.where(np.all(mask == label, axis=-1))[:2]] = iilabel_mask = label_mask.astype(int)return label_maskdef get_cityscapes_labels():return np.array([[128, 64, 128],[244, 35, 232],[70, 70, 70],[102, 102, 156],[190, 153, 153],[153, 153, 153],[250, 170, 30],[220, 220, 0],[107, 142, 35],[152, 251, 152],[0, 130, 180],[220, 20, 60],[255, 0, 0],[0, 0, 142],[0, 0, 70],[0, 60, 100],[0, 80, 100],[0, 0, 230],[119, 11, 32]])def get_pascal_labels():"""Load the mapping that associates pascal classes with label colorsReturns:np.ndarray with dimensions (21, 3)"""return np.asarray([[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], [192, 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]])
#添加get_Abyssinian_labels()函数
def get_Abyssinian_labels():return np.asarray([[0, 0, 0], [128, 0, 0]])

7、修改__init__.py

from dataloaders.datasets import cityscapes, coco, combine_dbs, pascal, sbd
from torch.utils.data import DataLoaderdef make_data_loader(args, **kwargs):if args.dataset == 'pascal':train_set = pascal.VOCSegmentation(args, split='train')val_set = pascal.VOCSegmentation(args, split='val')if args.use_sbd:sbd_train = sbd.SBDSegmentation(args, split=['train', 'val'])train_set = combine_dbs.CombineDBs([train_set, sbd_train], excluded=[val_set])num_class = train_set.NUM_CLASSEStrain_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True, **kwargs)val_loader = DataLoader(val_set, batch_size=args.batch_size, shuffle=False, **kwargs)test_loader = Nonereturn train_loader, val_loader, test_loader, num_classelif args.dataset == 'cityscapes':train_set = cityscapes.CityscapesSegmentation(args, split='train')val_set = cityscapes.CityscapesSegmentation(args, split='val')test_set = cityscapes.CityscapesSegmentation(args, split='test')num_class = train_set.NUM_CLASSEStrain_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True, **kwargs)val_loader = DataLoader(val_set, batch_size=args.batch_size, shuffle=False, **kwargs)test_loader = DataLoader(test_set, batch_size=args.batch_size, shuffle=False, **kwargs)return train_loader, val_loader, test_loader, num_classelif args.dataset == 'coco':train_set = coco.COCOSegmentation(args, split='train')val_set = coco.COCOSegmentation(args, split='val')num_class = train_set.NUM_CLASSEStrain_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True, **kwargs)val_loader = DataLoader(val_set, batch_size=args.batch_size, shuffle=False, **kwargs)test_loader = Nonereturn train_loader, val_loader, test_loader, num_class#添加Abyssinianelif args.dataset == 'Abyssinian':train_set = pascal.VOCSegmentation(args, split='train')val_set = pascal.VOCSegmentation(args, split='val')'''if args.use_sbd:sbd_train = sbd.SBDSegmentation(args, split=['train', 'val'])train_set = combine_dbs.CombineDBs([train_set, sbd_train], excluded=[val_set])'''num_class = train_set.NUM_CLASSEStrain_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True, **kwargs)val_loader = DataLoader(val_set, batch_size=args.batch_size, shuffle=False, **kwargs)test_loader = Nonereturn train_loader, val_loader, test_loader, num_classelse:raise NotImplementedError

8、修改train.py
在这个程序中，作者提供了四种的backbone，分别是：resnet、xception、drn、mobilenet，并且提供了预训练权重，在训练开始时会自动下载，可以根据自己电脑的配置进行选择，如果配置稍微差一点的话建议使用mobilenet作为backbone，在训练开始之前，打开train文件，在这里把自己的类别添加上去：
main函数的部分程序：

def main():parser = argparse.ArgumentParser(description="PyTorch DeeplabV3Plus Training")parser.add_argument('--backbone', type=str, default='resnet',choices=['resnet', 'xception', 'drn', 'mobilenet'],help='backbone name (default: resnet)')parser.add_argument('--out-stride', type=int, default=16,help='network output stride (default: 8)')#添加Abyssinianparser.add_argument('--dataset', type=str, default='pascal',choices=['pascal', 'coco', 'cityscapes'，'Abyssinian'],help='dataset name (default: pascal)')parser.add_argument('--use-sbd', action='store_true', default=True,help='whether to use SBD dataset (default: True)'

9、训练
在终端运行命令：

python train.py --backbone mobilenet --lr 0.007 --workers 1 --epochs 50 --batch-size 8 --gpu-ids 0 --checkname deeplab-mobilenet --dataset Abyssinian

每个轮次训练完后都会对验证集进行评估，会按照mIOU最优结果进行模型的保存，最后的模型的训练结果会保存在run文件夹中：

10、测试
添加测试代码：

#
# demo.py
#
import argparse
import os
import numpy as np
import timefrom modeling.deeplab import *
from dataloaders import custom_transforms as tr
from PIL import Image
from torchvision import transforms
from dataloaders.utils import  *
from torchvision.utils import make_grid, save_imagedef main():parser = argparse.ArgumentParser(description="PyTorch DeeplabV3Plus Training")parser.add_argument('--in-path', type=str, required=True, help='image to test')# parser.add_argument('--out-path', type=str, required=True, help='mask image to save')parser.add_argument('--backbone', type=str, default='resnet',choices=['resnet', 'xception', 'drn', 'mobilenet'],help='backbone name (default: resnet)')parser.add_argument('--ckpt', type=str, default='deeplab-resnet.pth',help='saved model')parser.add_argument('--out-stride', type=int, default=16,help='network output stride (default: 8)')parser.add_argument('--no-cuda', action='store_true', default=False,help='disables CUDA training')parser.add_argument('--gpu-ids', type=str, default='0',help='use which gpu to train, must be a \comma-separated list of integers only (default=0)')parser.add_argument('--dataset', type=str, default='pascal',choices=['pascal', 'coco', 'cityscapes','invoice'],help='dataset name (default: pascal)')parser.add_argument('--crop-size', type=int, default=513,help='crop image size')parser.add_argument('--num_classes', type=int, default=2,help='crop image size')parser.add_argument('--sync-bn', type=bool, default=None,help='whether to use sync bn (default: auto)')parser.add_argument('--freeze-bn', type=bool, default=False,help='whether to freeze bn parameters (default: False)')args = parser.parse_args()args.cuda = not args.no_cuda and torch.cuda.is_available()if args.cuda:try:args.gpu_ids = [int(s) for s in args.gpu_ids.split(',')]except ValueError:raise ValueError('Argument --gpu_ids must be a comma-separated list of integers only')if args.sync_bn is None:if args.cuda and len(args.gpu_ids) > 1:args.sync_bn = Trueelse:args.sync_bn = Falsemodel_s_time = time.time()model = DeepLab(num_classes=args.num_classes,backbone=args.backbone,output_stride=args.out_stride,sync_bn=args.sync_bn,freeze_bn=args.freeze_bn)ckpt = torch.load(args.ckpt, map_location='cpu')model.load_state_dict(ckpt['state_dict'])model = model.cuda()model_u_time = time.time()model_load_time = model_u_time-model_s_timeprint("model load time is {}".format(model_load_time))composed_transforms = transforms.Compose([tr.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),tr.ToTensor()])for name in os.listdir(args.in_path):s_time = time.time()image = Image.open(args.in_path+"/"+name).convert('RGB')# image = Image.open(args.in_path).convert('RGB')target = Image.open(args.in_path+"/"+name).convert('L')sample = {'image': image, 'label': target}tensor_in = composed_transforms(sample)['image'].unsqueeze(0)model.eval()if args.cuda:tensor_in = tensor_in.cuda()with torch.no_grad():output = model(tensor_in)grid_image = make_grid(decode_seg_map_sequence(torch.max(output[:3], 1)[1].detach().cpu().numpy()),3, normalize=False, range=(0, 255))save_image(grid_image,args.in_path+"/"+"{}_mask.png".format(name[0:-4]))u_time = time.time()img_time = u_time-s_timeprint("image:{} time: {} ".format(name,img_time))# save_image(grid_image, args.out_path)# print("type(grid) is: ", type(grid_image))# print("grid_image.shape is: ", grid_image.shape)print("image save in in_path.")
if __name__ == "__main__":main()# python demo.py --in-path your_file --out-path your_dst_file

在终端运行：

python demo.py --in-path  /home/yuxin/pytorch-deeplab-xception-master/test_pic --ckpt run/Abyssinian/deeplab-mobilenet/model_best.pth.tar --backbone mobilenet

效果如下图：

参考：
https://blog.csdn.net/fanxuelian/article/details/85145558
https://www.jianshu.com/p/026c5d78d3b1
https://blog.csdn.net/qq_39056987/article/details/106455828
https://blog.csdn.net/weixin_41919571/article/details/107906066 等