图像语义分割实战：TensorFlow Deeplabv3+ 训练自己数据集

---

前言

本文是为了记录deeplabv3+训练过程。

---

一、环境配置

我的环境：

ubuntu 16.04
anaconda3
python 3.5
tensorflow-gpu 1.10.0

anaconda3的安装可以上网查。

二、训练过程

1.引入库

首先clone官方提供的tensorflow/models文件。

git clone https://github.com/tensorflow/models.git

如果下载速度慢的话，可以在码云上下载。

2.数据集准备

转换为 VOC 格式的数据集

标注数据，制作符合要求的mask图像。利用图像对应的json文件，将数据转换成voc格式，方便后续进一步转换成deeplab训练所需的灰度图格式。

将labelme项目下载到本地：

git clone https://github.com/wkentaro/labelme.git

找到目录/labelme/examples/semantic_segmentation，里面有一个进行转换的完整示例，对照着示例，将自己的数据（原始图片和对应json标注）放入data_annotated文件夹，制作自己的labels.txt，拷贝labelme2voc.py文件不需改动，如下：

在当前目录下打开终端执行以下命令：

python labelme2voc.py data_annotated data_dataset_voc --labels labels.txt

会生成 data_dataset_voc 文件夹，里面包含：

JPEGImages文件夹中是原始图片。
SegmentationClassPNG文件夹中是生成的label图片。

Convert to 灰度图

deeplab使用单通道的标注图，即灰度图，并且类别的像素标记应该是0,1,2,3…n（共计n+1个类别，包含1个背景类和n个目标类）。执行remove_gt_colormap.py将mask转换成需要的格式，路径根据需要修改。

# from models/research/deeplab/datasets
python remove_gt_colormap.py \
  --original_gt_folder="/media/dell/2T/test/testlabel" \
  --output_dir="/media/dell/2T/test/mask"

original_gt_folder：原始标签图文件夹。
output_dir：要输出的标签图文件夹的位置。

Convert to tfrecord

制作tfrecord之前，需要将数据集分类成训练/测试/验证集。

数据集目录结构如下：

• data
  • image
  • mask
  • index
    • train.txt
    • trainval.txt
    • val.txt
  • tfrecord

iamge：存放所有的输入图片，包括训练、测试、验证集的图片。
mask：存放所有的labele（灰度）图片，和输入图片（即iamge）是一一对应的，文件名相同。
tfrecord：存放的是tfrecord格式的数据。
train.txt：所有训练集的文件名称（不包括后缀）
trainval.txt：所有验证集的文件名称（不包括后缀）
val.txt：所有测试集的文件名称（不包括后缀）

根据index下的txt文件运行build_voc2012_data.py转换成tfrecord格式，终端下执行以下代码：

# from /../models/research/deeplab/datasets/
python ./build_voc2012_data.py \
  --image_folder="/home/dell/models/research/deeplab/data/image" \
  --semantic_segmentation_folder="/home/dell/models/research/deeplab/data/mask" \
  --list_folder="/home/dell/models/research/deeplab/data/index" \
  --image_format="png" \
  --output_dir="/home/dell/models/research/deeplab/data/tfrecord"

image_folder ：数据集image的文件目录地址
semantic_segmentation_folder：数据集中mask的文件目录地址
list_folder : 将数据集分类成训练集、验证集等的指示目录index的文件目录
image_format : 输入图片数据的格式，我的数据集是png格式
output_dir：制作的TFRecord存放的目录地址

3.训练前代码准备

修改 deeplab/datasets/data_generator.py
在100行左右添加自己的数据集描述：

_MYDATA = DatasetDescriptor(
    splits_to_sizes={
        'train':150 ,  # num of samples in images/training 训练集
	    'trainval':168, 
        'val': 18,  # num of samples in images/validation 测试集
    },
    num_classes=3, #因为有两个标注的类，加上背景background类，所以共有3个类（根据标注的类+1）
    ignore_label=255, #背景类rgb，这里为黑色
)

然后在代码110行左右注册数据集：

_DATASETS_INFORMATION = {
    'cityscapes': _CITYSCAPES_INFORMATION,
    'pascal_voc_seg': _PASCAL_VOC_SEG_INFORMATION,
    'ade20k': _ADE20K_INFORMATION,
    'mydata':_MYDATA, #这里就是你上一步创建的数据类
}

修改train_utils.py文件
在train_utils.py中，先将大概209行的关于exclude_list的设置修改，作用是在使用预训练权重时候，不加载该logit层：

# Variables that will not be restored.
exclude_list = ['global_step','logits']
if not initialize_last_layer:
exclude_list.extend(last_layers)

4.主要训练参数

训练文件train.py和common.py文件中包含了训练分割网络所需要的所有参数。

• model_variant：Deeplab模型变量，可选值可见core/feature_extractor.py。

  • 当使用mobilenet_v2时，设置变量strous_rates=decoder_output_stride=None；
  • 当使用xception_65或resnet_v1时，设置strous_rates=[6,12,18](output stride 16)， decoder_output_stride=4。

• label_weights：此变量可以设置标签的权重值，当数据集中出现类别不均衡时，可通过此变量来指定每个类别标签的权重值，如label_weights=[0.1, 0.5]意味着标签0的权重是0.1， 标签1的权重是0.5。如果该值为None，则所有的标签具有相同的权重1.0。
  提示：训练中我没有修改这个变量，下次训练时会尝试设置。

• train_logdir：存放checkpoint和logs的路径。

• log_steps：该值表示每隔多少步输出日志信息。

• save_interval_secs：该值表示以秒为单位，每隔多长时间保存一次模型文件到硬盘。

• optimizer：优化器，可选值[‘momentum’, ‘adam’]。默认为momentum。

• learning_policy：学习率策略，可选值[‘poly’, ‘step’]。

• base_learning_rate：基础学习率，默认值0.0001。

• training_number_of_steps：模型训练的迭代次数。

• train_batch_size：模型训练的批处理图像数量。

• train_crop_size：模型训练时所使用的图像尺寸，默认’513, 513’。

• tf_initial_checkpoint：预训练模型。

• initialize_last_layer：是否初始化最后一层。

• last_layers_contain_logits_only：是否只考虑逻辑层作为最后一层。

• fine_tune_batch_norm：是否微调batch norm参数。

• atrous_rates：默认值[6, 12, 18]。

• output_stride：默认值16，输入和输出空间分辨率的比值

  • 对于xception_65， 如果output_stride=8，则使用atrous_rates=[12, 24, 36]
  • 如果output_stride=16，则atrous_rates=[6, 12, 18]
  • 对于mobilenet_v2，使用None
  • 注意：在训练和验证阶段可以使用不同的atrous_rates和output_stride。

• dataset：所使用的分割数据集，此处与数据集注册时的名称一致。

• train_split：使用哪个数据集来训练，可选值即数据集注册时的值，如train, trainval。

• dataset_dir：数据集存放的路径。
  针对训练参数，下面几点需要重点注意：

1.关于是否加载预训练网络的权重问题
如果要在其他数据集上微调该网络，需要关注以下几个参数：

1. 使用预训练网络的权重，设置initialize_last_layer=True
2. 只使用网络的backbone，设置initialize_last_layer=False和last_layers_contain_logits_only=False
3. 使用所有的预训练权重，除了logits，因为如果是自己的数据集，对应的classes不同（这个我们前面已经设置不加载logits）设置initialize_last_layer=False和last_layers_contain_logits_only=True
4. 由于我的数据集分类与默认类别数不同，因此采取的参数值是：

--initialize_last_layer=false
--last_layers_contain_logits_only=true

2.如果资源有限，想要训练自己数据集的几条建议：

1. 设置output_stride=16或者甚至32（同时需要修改atrous_rates变量，例如，对于output_stride=32，atrous_rates=[3, 6, 9]）
2. 尽可能多的使用GPU，更改num_clone标志，并将train_batch_size设置的尽可能大
3. 调整train_crop_size，可以将它设置的更小一些，例如513x513（甚至321x321），这样就可以使用更大的batch_size
4. 使用较小的网络主干，如mobilenet_v2

3.关于是否微调batch_norm

1. 当训练使用的批处理大小train_batch_size大于12（最好大于16）时，设置fine_tune_batch_norm=True。否则，设置fine_tune_batch_norm=False。

5.预训练模型

选择预训练模型，根据自己的情况进行下载，下载地址：https://github.com/tensorflow/models/blob/master/research/deeplab/g3doc/model_zoo.md
我选择的是xception_71。

6.测试model_test.py

测试一下环境配置是否成功。

添加依赖库到PYTHONPATH，在目录/home/user/models/research/下：

#From /home/user/models/research/
export PYTHONPATH=$PYTHONPATH:`pwd`:`pwd`/slim
source ~/.bashrc

调用model_test.py测试：

# From /home/user/models/research/
python deeplab/model_test.py

7.训练

train.py：训练代码文件，训练时，需要指定提供的训练参数。

# from /../models/research/
python deeplab/train.py \
    --logtostderr \
    --training_number_of_steps=80000 \
    --train_split="train" \
    --model_variant="xception_71" \
    --atrous_rates=6 \
    --atrous_rates=12 \
    --atrous_rates=18 \
    --output_stride=16 \
    --decoder_output_stride=4 \
    --train_crop_size="321,321"\
    --train_batch_size=8 \
    --fine_tune_batch_norm = False \
    --base_learning_rate=0.01 \
    --dataset="mydata" \
    --tf_initial_checkpoint='/home/dell/models/research/deeplab/backbone/xception_71/model.ckpt' \
    --train_logdir='/home/dell/models/research/deeplab/exp/mydata_train/train' \
    --dataset_dir='/home/dell/models/research/deeplab/data/tfrecord'

8.可视化测试

vis.py：可视化代码。

# from /root/models/research/
python deeplab/vis.py \
    --logtostderr \
    --vis_split="val" \
    --model_variant="xception_71" \
    --atrous_rates=6 \
    --atrous_rates=12 \
    --atrous_rates=18 \
    --output_stride=16 \
    --decoder_output_stride=4 \
    --vis_crop_size="512,512" \
    --dataset="mydata" \
    --colormap_type="pascal" \
    --checkpoint_dir='/home/dell/models/research/deeplab/exp/mydata_train/train/' \
    --vis_logdir='/home/dell/models/research/deeplab/exp/mydata_train/vis/' \
    --dataset_dir='/home/dell/models/research/deeplab/data/tfrecord/' \
    --max_number_of_iterations=1

9.验证

eval.py：验证代码，输出mIOU，用来评估模型的好坏。

# from /root/models/research/
python deeplab/eval.py \
    --logtostderr \
    --eval_split="val" \
    --model_variant="xception_71" \
    --atrous_rates=6 \
    --atrous_rates=12 \
    --atrous_rates=18 \
    --output_stride=16 \
    --decoder_output_stride=4 \
    --eval_crop_size="512,512" \
    --dataset="mydata" \
    --checkpoint_dir='/home/dell/models/research/deeplab/exp/mydata_train/train/' \
    --eval_logdir='/home/dell/models/research/deeplab/exp/mydata_train/eval/' \
    --dataset_dir='/home/dell/models/research/deeplab/data/tfrecord/' \
    --max_number_of_iterations=1

10.查看日志

使用Tensorboard检查培训和评估工作的进展。

tensorboard --logdir=${PATH_TO_LOG_DIRECTORY}
# 文中log地址
tensorboard --logdir="./train_logs"

可以使用以下命令运行:

#查看eval日志
tensorboard --logdir=/home/dell/models/research/deeplab/exp/mydata_train/eval --port 6007

#查看训练日志
tensorboard --logdir=/home/dell/models/research/deeplab/exp/mydata_train/train

11.导出模型

在训练过程中，会保存模型文件到硬盘，如下：

代码中提供了一个脚本（export_model.py）可以将checkpoint转换为.pb格式。
在./models/research/下创建一个脚本export_model.sh用来执行export_model.py，内容为export_model.py的主要修改参数，代码如下：

export PYTHONPATH=$PYTHONPATH:`pwd`:`pwd`/slim
python deeplab/export_model.py \
    --logtostderr \
    --checkpoint_path="/media/dell/2T/models/research/deeplab/exp/mydata_train/train/model.ckpt-$1" \
    --export_path="/media/dell/2T/models/research/deeplab/exp/mydata_train/export/inference_graph-$1.pb" \
    --model_variant="xception_71" \
    --atrous_rates=6 \
    --atrous_rates=12 \
    --atrous_rates=18 \
    --output_stride=16 \
    --decoder_output_stride=4 \
    --num_classes=2 \
    --crop_size=3000 \
    --crop_size=3000 \
    --inference_scales=1.0
# checkpoint_path：训练保存的检查点文件
# export_path：模型导出路径
# num_classes：分类类别
# crop_size：图像尺寸，默认为[513, 513]，可以修改大小，这个影响测试图像输出的图像尺寸。但是尺寸过大时导出模型可能会报错，好像和显存大小有关，而且尺寸越大的话，测试图像需要的时间会越长。
# atrous_rates和output_stride可以和训练的时候不一样。我的配置是一样的。
# inference_scales：多尺度推理，默认[1.0]。更改为 [0.5, 0.75, 1.0, 1.25, 1.5, 1.75] 进行多尺度推理。

运行export_model.sh导出模型。

# from /../models/research/
sh export_model.sh 80000
# 数字表示模型名称后面的数字

生成的.pb文件如下：

三、测试

可以自己写测试代码，下面的代码为其他文章中直接搬来可作为参考，自行修改和精简，示例如下：

#!/usr/bin/env python
# coding: utf-8
 
import os
from io import BytesIO
import tarfile
import tempfile
from six.moves import urllib
from matplotlib import gridspec
from matplotlib import pyplot as plt
import numpy as np
from PIL import Image
import tensorflow as tf
import scipy
 
LABEL_NAMES = np.asarray(["background", "class1", "class2"])
 
 
class DeepLabModel(object):
    """Class to load deeplab model and run inference."""
 
    INPUT_TENSOR_NAME = "ImageTensor:0"
    OUTPUT_TENSOR_NAME = "SemanticPredictions:0"
    INPUT_SIZE = 321
    FROZEN_GRAPH_NAME = "frozen_inference_graph"
 
    def __init__(self, modelname):
        """Creates and loads pretrained deeplab model."""
        self.graph = tf.Graph()
        graph_def = None
 
        with open(modelname, "rb") as fd:
            graph_def = tf.GraphDef.FromString(fd.read())
 
        if graph_def is None:
            raise RuntimeError("Cannot find inference graph in tar archive.")
 
        with self.graph.as_default():
            tf.import_graph_def(graph_def, name="")
 
        self.sess = tf.Session(graph=self.graph)
 
    def run(self, image):
        """Runs inference on a single image.
        Args:
        image: A PIL.Image object, raw input image.
        Returns:
        resized_image: RGB image resized from original input image.
        seg_map: Segmentation map of `resized_image`.
        """
        width, height = image.size
        resize_ratio = 1.0 * self.INPUT_SIZE / max(width, height)
        target_size = (int(resize_ratio * width), int(resize_ratio * height))
        resized_image = image.convert("RGB").resize(target_size, Image.ANTIALIAS)
        batch_seg_map = self.sess.run(
            self.OUTPUT_TENSOR_NAME,
            feed_dict={self.INPUT_TENSOR_NAME: [np.asarray(resized_image)]},
        )
        seg_map = batch_seg_map[0]
        return resized_image, seg_map
 
 
def create_pascal_label_colormap():
    """Creates a label colormap used in PASCAL VOC segmentation benchmark.
    Returns:
        A Colormap for visualizing segmentation results.
    """
    colormap = np.zeros((256, 3), dtype=int)
    ind = np.arange(256, dtype=int)
 
    for shift in reversed(range(8)):
        for channel in range(3):
            colormap[:, channel] |= ((ind >> channel) & 1) << shift
        ind >>= 3
 
    return colormap
 
 
# 从 label 到 color_image
def label_to_color_image(label):
    """Adds color defined by the dataset colormap to the label.
    Args:
        label: A 2D array with integer type, storing the segmentation label.
    Returns:
        result: A 2D array with floating type. The element of the array
        is the color indexed by the corresponding element in the input label
        to the PASCAL color map.
    Raises:
        ValueError: If label is not of rank 2 or its value is larger than color
        map maximum entry.
    """
    if label.ndim != 2:
        raise ValueError("Expect 2-D input label")
 
    colormap = create_pascal_label_colormap()
 
    if np.max(label) >= len(colormap):
        raise ValueError("label value too large.")
 
    return colormap[label]
 
 
# 分割结果可视化
def vis_segmentation(image, seg_map, name):
    """Visualizes input image, segmentation map and overlay view."""
    plt.figure(figsize=(15, 5))
    grid_spec = gridspec.GridSpec(1, 4, width_ratios=[6, 6, 6, 1])
 
    plt.subplot(grid_spec[0])
    plt.imshow(image)
    plt.axis("off")
    plt.title("input image")
 
    plt.subplot(grid_spec[1])
    seg_image = label_to_color_image(seg_map).astype(np.uint8)
    plt.imshow(seg_image)
    plt.axis("off")
    plt.title("segmentation map")
 
    plt.subplot(grid_spec[2])
    plt.imshow(image)
    plt.imshow(seg_image, alpha=0.7)
    plt.axis("off")
    plt.title("segmentation overlay")
 
    unique_labels = np.unique(seg_map)
    ax = plt.subplot(grid_spec[3])
    plt.imshow(FULL_COLOR_MAP[unique_labels].astype(np.uint8), interpolation="nearest")
    ax.yaxis.tick_right()
    plt.yticks(range(len(unique_labels)), LABEL_NAMES[unique_labels])
    plt.xticks([], [])
    ax.tick_params(width=0.0)
    plt.grid("off")
 
    plt.savefig("./seg_map_result/" + name + ".png")
    # plt.show()
 
 
FULL_LABEL_MAP = np.arange(len(LABEL_NAMES)).reshape(len(LABEL_NAMES), 1)
FULL_COLOR_MAP = label_to_color_image(FULL_LABEL_MAP)
 
 
def main_test(filepath):
    # 加载模型
    modelname = "./datasets/quekou/export/inference_graph-80000.pb"
    MODEL = DeepLabModel(modelname)
    print("model loaded successfully!")
 
    filelist = os.listdir(filepath)
    for item in filelist:
        print("process image of ", item)
        name = item.split(".jpg", 1)[0]
        original_im = Image.open(filepath + item)
        resized_im, seg_map = MODEL.run(original_im)
 
        # 分割结果拼接
        vis_segmentation(resized_im, seg_map, name)
 
        # 单独保存分割结果
        # seg_map_name = name + '_seg.png'
        # resized_im_name = name + '_in.png'
        # path = './seg_map_result/'
        # scipy.misc.imsave(path + resized_im_name,resized_im)
        # scipy.misc.imsave(path + seg_map_name,seg_map)
 
 
if __name__ == "__main__":
    filepath = "./datasets/quekou/dataset/JPEGImages/"
    main_test(filepath)

参考文章：

原文链接：https://blog.csdn.net/malvas/article/details/90776327
原文链接：https://blog.csdn.net/ling620/article/details/105635780
原文链接：https://blog.csdn.net/zong596568821xp/article/details/83350820

---

总结

以上就是今天要讲的内容，本文记录了tensorflow下利用deeplabv3+对自己的数据集进行训练和测试。训练过程中可能还有其他问题，日后会继续补充。