图像分割：Tensorflow Deeplabv3+训练人像分割数据集

冰蓝儿
4,116 阅读3分钟

本文主要介绍如何使用自己的数据集训练DeepLabv3+分割算法，代码使用的是官方源码。

1、代码简介

当前使用TensorFlow版本的官方源码，选择它的原因是因为代码中的内容比较全面，除了代码实现以外，还提供了许多文档帮助理解与使用，同时还提供了模型转换的代码实现。

代码地址： 【github】models/research/deeplab at master · tensorflow/models

接下来，先对这个代码仓库进行一下简单的介绍，因为自己在使用该代码仓库的时候只关心训练代码的实现，而忽略的其他的内容，走了不少弯路，到后面才发现我想要的内容，仓库里面早有（==）。

在当前的实现中，我们支持采用以下网络主干：

  • MobileNetv2和MobileNetv3：一个为移动设备设计的快速网络结构
  • Xception:用于服务器端部署的强大网络结构
  • ResNet-v1-{50, 101}:我们提供原始的ResNet-v1及其“ beta”变体，其中对“ stem”进行了修改以进行语义分割。
  • PNASNet: 一个通过神经体系结构搜索发现的强大网络结构。
  • Auto-Deeplab（代码中叫做HNASNet）：通过神经体系结构搜索找到的特定于细分的网络主干。

该目录包含TensorFlow 实现。我们提供的代码使用户可以训练模型，根据mIOU（平均交叉点求和）评估结果以及可视化细分结果。我们以PASCAL VOC 2012和Cityscapes语义分割基准为例。

代码中几个重要文件：

  • datasets/：该文件夹下包含对于训练数据集的处理代码，主要针对 PASCAL VOC 2012和Cityscapes数据集的处理。
  • g3doc/：该文件夹下包含多个Markdown文件，非常有用，如何安装，常见问题FAQ等。
  • deeplab_demo.ipynb：该文件中给出了如果对一张图像进行语义分割并显示结果的Demo。
  • export_model.py：该文件提供了将训练的checkpoint模型转为.pb文件的代码实现。
  • train.py：训练代码文件，训练时，需要指定提供的训练参数。
  • eval.py：验证代码，输出mIOU，用来评估模型的好坏。
  • vis.py：可视化代码。

2、安装

Deeplab依赖的库有：

  • Numpy
  • Pillow 1.0
  • tf Slim (which is included in the "tensorflow/models/research/" checkout)
  • Jupyter notebook
  • Matplotlib
  • Tensorflow

2.1 添加库到PYTHONPATH

本地运行的时候，tensorflow/models/research/目录应该追加到PYTHONPATH中，如下：

# From tensorflow/models/research/
export PYTHONPATH=$PYTHONPATH:`pwd`:`pwd`/slim

# [Optional] for panoptic evaluation, you might need panopticapi:
# https://github.com/cocodataset/panopticapi
# Please clone it to a local directory ${PANOPTICAPI_DIR}
touch ${PANOPTICAPI_DIR}/panopticapi/__init__.py
export PYTHONPATH=$PYTHONPATH:${PANOPTICAPI_DIR}/panopticapi

注意：此命令需要在您启动的每个新终端上运行。如果希望避免手动运行此命令，可以将它作为新行添加到〜/ .bashrc文件的末尾。

2.2 测试是否安装成功

通过运行model_test.py快速测试：

# From tensorflow/models/research/
python deeplab/model_test.py

在PASCAL VOC 2012数据集上快速运行所有代码：

# From tensorflow/models/research/deeplab
sh local_test.sh

3、数据集准备

最终目标： 生成TFRecord格式的数据

数据集目录结构如下：

+dataset #数据集名称
	+image
	+mask
	+index
		- train.txt
		- trainval.txt
		- val.txt
	+tfrecord
  • image： 原图图像，RGB彩色图像
  • mask：像素值为类别标签的mask图像，单通道，与原图的名称一致，后缀为.jpg和.png都可以，只要在代码中读取一致即可。VOC数据集默认原图是.jpg，mask图像为.png。
  • index：存放图像文件名的txt文件（不加后缀）
  • tfrecord：存放转为tfrecord格式的图像数据

数据集制作流程：

  1. 标注数据，制作符合要求的mask图像
  2. 将数据集分割为训练集、验证集和测试集
  3. 生成TFRecord格式的数据集

3.1 标注数据

训练集数据包含两部分，一是原图，二是对应分类的标注值（本文中称为mask图像）。

mask图像的值是如何设置的？ 根据图像分割的分类个数来制作原图对应的mask图像。假如一共有N个类别（背景作为一类），则mask图像的值的范围是[0~N)。0值作为背景值，其他分割类别的值依次设置为1, 2, ..., N-1。

注意：

  • ignore_label：从字面意思来讲是忽略的标签，即ignore_label是指没有做标注的像素，即不需进行预测的像素值，因此，它不参与loss值的计算，在mask图像中将其值记为255。
  • mask图像是单通道的灰度图像。
  • mask图像的格式没有限定，但所有的mask图像采用同一种图像格式，方便数据读取。

小总结 mask图像的值分为三类：

  1. 背景：用0表示
  2. 分类类别：使用1, 2, ....., N-1表示
  3. ignore_label值：用255表示

如果分割的类别较少，则生成的mask图像看上去是一片黑，因为分类的值都较小，在0~255的范围内不容易显示出来。

3.2 分割数据集

这部分就是将准备的数据集进行分割，分为训练集、验证集、测试集。 无需将具体的图像文件分到三个文件夹中，只需要建立图像的索引文件即可，通过添加相应的路径+文件名即可获取到具体的图像。

假设原图像和mask图像的存放路径如下：

  • 原图：./dataset/images
  • mask图像：./dataset/mask：此处存放的是2.1小节要求格式

原图与mask图像是一一对应的，包括图像尺寸，图像名（后缀可以不同）

索引文件存放路径：./dataset/index，该路径下生成：

  • train.txt
  • trainval.txt
  • val.txt

索引文件中，只需记录文件名（不加后缀），这取决于代码中数据集加载的方式。

3.3 将数据打包为TFRecord格式

TFRecord是谷歌推荐的一种二进制文件格式，理论上它可以保存任何格式的信息。 TFRecord内部使用了“Protocol Buffer”二进制数据编码方案，它只占用一个内存块，只需要一次性加载一个二进制文件的方式即可，简单，快速，尤其对大型训练数据很友好。而且当我们的训练数据量比较大的时候，可以将数据分成多个TFRecord文件，来提高处理效率。

那么，如何将数据生成TFRecord格式呢？

在此，我们可以借助 项目代码中./datasets/build_voc2012_data.py文件来实现。给文件是VOC2012数据集处理的代码，我们只需修改一下输入参数即可。

参数：

  • image_folder：原图文件夹名称，./dataset/image
  • semantic_segmentation_folder：分割文件夹名称, ./dataset/mask
  • list_folder：索引文件夹名称，./dataset/index
  • output_dir：输出路径，即生成的tfrecord文件所在位置，./dataset/tfrecord

运行命令：

python ./datasets/build_voc2012_data.py --image_folder=./dataset/image
										--semantic_segmentation_folder=./dataset/mask
										--list_folder=./dataset/index
										--output_dir=./dataset/tfrecord

生成的文件如下：

在这里插入图片描述
注意： 可在代码中调节参数_NUM_SHARDS （默认为4），改变数据分块的数目。（一些文件系统有最大单个文件大小的限制，如果数据集非常大，增加_NUM_SHARDS 可减小单个文件的大小）

该文件的核心代码如下：

# dataset_split指的是train.txt, val.txt等
dataset = os.path.basename(dataset_split)[:-4]
filenames = [x.strip('\n') for x in open(dataset_split, 'r')] # 文件名列表

# 输出tfrecord文件名
output_filename = os.path.join(
            FLAGS.output_dir,
            '%s-%05d-of-%05d.tfrecord' % (dataset, shard_id, _NUM_SHARDS))
with tf.python_io.TFRecordWriter(output_filename) as tfrecord_writer:
	for i in range(start_idx, end_idx): 
		image_filename = os.path.join(iamge_folder, filenames[i]+'.'+image_format)# 原图路径
		image_data = tf.gfile.GFile(image_filename, 'rb').read() #读取原图文件
	    height, width = image_reader.read_image_dims(image_data)
	    
	    seg_filename = os.path.join(semantic_segmentation_folder,
                    filenames[i] + '.' + label_format) # mask图像路径
	    seg_data = tf.gfile.GFile(seg_filename, 'rb').read() # 读取分割图像
	    seg_height, seg_width = label_reader.read_image_dims(seg_data)
	    
	    # 判断原图与mask图像尺寸是否匹配
	    if height != seg_height or width != seg_width:
	        raise RuntimeError('Shape mismatched between image and label.')
	    # Convert to tf example.
	    example = build_data.image_seg_to_tfexample(
	        image_data, filenames[i], height, width, seg_data)
	    tfrecord_writer.write(example.SerializeToString())

至此，数据集的制作部分已经完成！！！

4、训练

4.1 代码修改

为了训练自己的数据集，需要修改以下几处文件：

1 datasets/data_generator.py：增加数据集的注册

该文件提供语义分割数据的包装器

在该文件中，可以看到PASCAL_VOC， CITYSCAPES以及ADE20K数据集的数据描述，如下：

_PASCAL_VOC_SEG_INFORMATION = DatasetDescriptor(
    splits_to_sizes={
        'train': 1464,
        'train_aug': 10582,
        'trainval': 2913,
        'val': 1449,
    },
    num_classes=21,
    ignore_label=255,
)

en，比着葫芦画瓢，增加我们自己数据集的描述信息，如下：

_PORTRAIT_INFORMATION = DatasetDescriptor(
    splits_to_sizes={
        'train': 17116,
        'trainval': 21395,
        'val': 4279,
    },
    num_classes=2,  # 类别数目，包括背景
    ignore_label=255,  # 忽略像素值
)

以人像分割任务为例，只有两类，即前景（人像）和背景（非人像）。

添加完描述信息后，需要将该数据集信息进行注册，如下：

_DATASETS_INFORMATION = {
    'cityscapes': _CITYSCAPES_INFORMATION,
    'pascal_voc_seg': _PASCAL_VOC_SEG_INFORMATION,
    'ade20k': _ADE20K_INFORMATION,
    'portrait_seg': _PORTRAIT_INFORMATION, #增加此句
}

注意：此处的数据集名称要与前面对应！

2 ./utils/train_utils.py修改

在函数get_model_init_fn中，修改为如下代码，增加logits层不加载预训练模型权重：

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

4.2 主要训练参数

训练文件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']。
  • 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
    • 注意：在训练和验证阶段可以使用不同的strous_rates和output_stride。
  • dataset：所使用的分割数据集，此处与数据集注册时的名称一致。
  • train_split：使用哪个数据集来训练，可选值即数据集注册时的值，如train, trainval。
  • dataset_dir：数据集存放的路径。

针对训练参数，下面几点需要重点注意：

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

    • 使用预训练网络的权重，设置initialize_last_layer=True
    • 只使用网络的backbone，设置initialize_last_layer=False和last_layers_contain_logits_only=False
    • 使用所有的预训练权重，除了logits，设置initialize_last_layer=False和last_layers_contain_logits_only=True

    由于我的数据集分类与默认类别数不同，因此采取的参数值是：

    --initialize_last_layer=false
--last_layers_contain_logits_only=true

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

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

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

4.3 预训练模型

模型链接具体可见：models/model_zoo.md at master · tensorflow/models

提供了在几个数据集上的预训练模型，包括(1) PASCAL VOC 2012， (2) Cityscapes， (3) ADE20K

未解压的目下包括：

  • 一个frozen inference graph（forzen_inference_graph.pb）。默认情况下，所有冻结推理图的输出步长为8，单个eval scale为1.0，没有左右翻转，除非另外指定。基于MobileNet-v2的模型不包括解码器模块。
  • 一个checkpoint(model.ckpt.data-00000-of-00001, model.ckpt.index)

还提供了在ImageNet预训练的checkpoints

未解压文件包括： 一个model checkpoint (model.ckpt.data-00000-of-00001, model.ckpt.index)

根据自己的情况进行下载！

4.4 训练模型

python train.py \
    --logtostderr \
    --training_number_of_steps=20000 \
    --train_split="train" \
    --model_variant="xception_65" \
    --train_crop_size="513,513" \
    --atrous_rates=6 \
    --atrous_rates=12 \
    --atrous_rates=18 \
    --output_stride=16 \
    --decoder_output_stride=4 \
    --train_batch_size=2 \
    --save_interval_secs=240 \
    --optimizer="momentum" \
    --leraning_policy="poly" \
    --fine_tune_batch_norm=false \
    --initialize_last_layer=false \
    --last_layers_contain_logits_only=true \
    --dataset="portrait_seg" \
    --tf_initial_checkpoint="./checkpoint/deeplabv3_pascal_trainval/model.ckpt" \
    --train_logdir="./train_logs" \
    --dataset_dir="./dataset/tfrecord"

4.5 验证模型

验证代码： ./eval.py

# From tensorflow/models/research/
python deeplab/eval.py \
    --logtostderr \
    --eval_split="val" \
    --model_variant="xception_65" \
    --atrous_rates=6 \
    --atrous_rates=12 \
    --atrous_rates=18 \
    --output_stride=16 \
    --decoder_output_stride=4 \
    --eval_crop_size="513,513" \
    --dataset="portrait_seg" \ # 数据集名称
    --checkpoint_dir=${PATH_TO_CHECKPOINT} \ # 预训练模型
    --eval_logdir=${PATH_TO_EVAL_DIR} \ 
    --dataset_dir="./dataset/tfrecord" # 数据集路径

得到的结果如下：

在这里插入图片描述

4.6 训练过程可视化

可以使用Tensorboard检查培训和评估工作的进展。如果使用推荐的目录结构，Tensorboard可以使用以下命令运行:

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

5、推理

5.1 模型导出

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

其形式是TensorFlow的checkpoint格式，代码中提供了一个脚本（export_model.py）可以将checkpoint转换为.pb格式。

export_model.py主要参数：

  • checkpoint_path：训练保存的检查点文件
  • export_path：模型导出路径
  • num_classes：分类类别
  • crop_size：图像尺寸，[513, 513]
  • atrous_rates：12, 24, 36
  • output_stride：8

生成的.pb文件如下：

在这里插入图片描述

5.2 单张图像上推理

class DeepLabModel(object):
    """class to load deeplab model and run inference"""

    INPUT_TENSOR_NAME = 'ImageTensor:0'
    OUTPUT_TENSOR_NAME='SemanticPredictions:0'
    INPUT_SIZE = 513
    FROZEN_GRAPH_NAME= 'frozen_inference_graph'

    def __init__(self, pretrained_weights):
        """Creates and loads pretrained deeplab model."""
        self.graph = tf.Graph()
        graph_def = None
        # Extract frozen graph from tar archive
        if pretrained_weights.endswith('.tar.gz'):
            tar_file = tarfile.open(pretrained_weights)
            for tar_info in tar_file.getmembers():
                if self.FROZEN_GRAPH_NAME in os.path.basename(tar_info.name):
                    file_handle = tar_file.extractfile(tar_info)
                    graph_def = tf.GraphDef.FromString(file_handle.read())
                    break
            tar_file.close()
        else:
            with open(pretrained_weights, '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='')

        gpu_options = tf.GPUOptions(allow_growth=True)
        config = tf.ConfigProto(gpu_options=gpu_options, log_device_placement=False)
        self.sess = tf.Session(graph=self.graph, config=config)

    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_iamge'.
        """
        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
        
        
if __name__ == '__main__':
    pretrained_weights = './train_logs/frozen_inference_graph_20000.pb'
    MODEL = DeepLabModel(pretrained_weights) # 加载模型
    
    img_name = 'test.jpg'
    img = Image.open(img_name)
    resized_im, seg_map = MODEL.run(original_im) #获取结果
    seg_map[seg_map==1]=255 #将人像的像素值置为255
    seg_map.save('output.jpg') # 保存mask结果图像

至此，整个训练过程就结束了！！！

avatar
文章
阅读
粉丝