TensorFlow DeepLabV3 训练自己的数据分割模型

RobinHoodltm9l 2019-10-19

展开全文

DeepLabV3+

deeplab v3+ 算是目前来说最先进的语义分割算法，尽管现在有精确到头发丝的分割方法：Soft Semantic Segmentation. 但谷歌开源了deeplabv3+，我们可以直接使用不同的backbone和数据集来训练我们自己的分割模型。如果你想做一个分割的应用，那这个方法再合适不过。

在开始之前，让我们看看我们能做到的分割效果：

当然出来cityspcapes数据集之外你也可以放许多其他的数据集。包括我自己的车道线数据集，分割效果也还不错，连左右车道线都能分割出来。接下来就教大家如何制作数据集吧。

数据集制作

首先这个数据集制作有个巨大的问题。假设你用labelme或者其他工具标注了你的数据，你的保存标注可能是polygon的点，也可能是mask。这里我推荐保存polygon，因为deeplab中使用的label是单通道以你的类别的id为像素值的标签，这也非常好理解。这里贴一个核心代码：

 1import cv2
 2import os
 3import numpy as np
 4import cv2
 5import matplotlib.pyplot as plt
 6
 7cls_label_map = {
 8    'background': 0,
 9    'a': 1,
10    's': 2,
11    'f': 3,
12    'd': 4,
13    'j': 5,
14    'k': 6,
15    'l': 7,
16    'g': 8,
17    'p': 9,
18    'q': 10,
19    'o': 11,
20    'w': 12,
21    'h': 13,
22    '3': 14,
23}
24
25
26def gen_mask(polygons, img):
27    mask = np.zeros([img.shape[0], img.shape[1]], np.int32)
28    for p in polygons:
29        # cv2.polylines(img, p, 1, 255)
30        # fill polygon
31        pts = np.array(p[:-1])
32        cv2.fillPoly(mask, pts, cls_label_map[p[-1]])
33    return mask

这样返回就是label的mask。这个mask实际上是黑色的。咋一看like this：

仔细看是可以隐约看到车道线的。并且其实这里的车道先不止一种类别。

关于label
关于label有几点需要注意的，像素值就是label的index，从我的map也能看的出来。
除此之外没了。另外，如果你的类别里面没有ignore_label, 那就直接是idx和0,0就是背景。如果有ignore_label就是255,相应的类别写进去，颜色值为255就是ignore了。

这只是第一步，接下来你的生成相应的tfrecords。这里我也提供一个核心代码，大家可以相应的进行一些更改。

  1import math
  2import os.path
  3import sys
  4import build_data
  5import tensorflow as tf
  6
  7FLAGS = tf.app.flags.FLAGS
  8
  9tf.app.flags.DEFINE_string('image_folder',
 10                           './VOCdevkit/VOC2012/JPEGImages',
 11                           'Folder containing images.')
 12
 13tf.app.flags.DEFINE_string(
 14    'semantic_segmentation_folder',
 15    './VOCdevkit/VOC2012/SegmentationClassRaw',
 16    'Folder containing semantic segmentation annotations.')
 17
 18tf.app.flags.DEFINE_string(
 19    'list_folder',
 20    './VOCdevkit/VOC2012/ImageSets/Segmentation',
 21    'Folder containing lists for training and validation')
 22
 23tf.app.flags.DEFINE_string(
 24    'output_dir',
 25    './tfrecord',
 26    'Path to save converted SSTable of TensorFlow examples.')
 27
 28
 29_NUM_SHARDS = 4
 30
 31def load_images_and_labels(d):
 32    all_dirs = [i for i in os.listdir(d) if i.startswith('lane_2018')]
 33    all_images = []
 34    all_labels = []
 35    for d in all_dirs:
 36        img_dir = os.path.join(d, 'images')
 37        label_dir = os.path.join(d, 'masks')
 38        all_images.extend([os.path.join(img_dir, i) for i in os.listdir(os.path.join(d, 'images')) if i.endswith('jpg') or i.endswith('jpeg')])
 39        all_labels.extend([os.path.join(img_dir, i) for i in os.listdir(os.path.join(d, 'masks')) if i.endswith('png')])
 40    return all_images, all_labels
 41
 42
 43def _convert_dataset(d):
 44    '''Converts the specified dataset split to TFRecord format.
 45
 46    Args:
 47      dataset_split: The dataset split (e.g., train, test).
 48
 49    Raises:
 50      RuntimeError: If loaded image and label have different shape.
 51    '''
 52    sys.stdout.write('Processing...')
 53
 54    all_image_files, all_label_files = load_images_and_labels(d)
 55
 56    num_images = len(all_image_files)
 57    num_per_shard = int(math.ceil(num_images / float(_NUM_SHARDS)))
 58
 59    image_reader = build_data.ImageReader('jpeg', channels=3)
 60    label_reader = build_data.ImageReader('png', channels=1)
 61
 62    dataset = 'train'
 63    for shard_id in range(_NUM_SHARDS):
 64        output_filename = os.path.join(
 65            FLAGS.output_dir,
 66            '%s-%05d-of-%05d.tfrecord' % (dataset, shard_id, _NUM_SHARDS))
 67        with tf.python_io.TFRecordWriter(output_filename) as tfrecord_writer:
 68            start_idx = shard_id * num_per_shard
 69            end_idx = min((shard_id + 1) * num_per_shard, num_images)
 70            for i in range(start_idx, end_idx):
 71                sys.stdout.write('\r>> Converting image %d/%d shard %d' % (
 72                    i + 1, len(filenames), shard_id))
 73                sys.stdout.flush()
 74                # Read the image.
 75                # image_filename = os.path.join(
 76                #     FLAGS.image_folder, filenames[i] + '.' + FLAGS.image_format)
 77                image_data = tf.gfile.FastGFile(all_image_files[i], 'rb').read()
 78                height, width = image_reader.read_image_dims(image_data)
 79                # Read the semantic segmentation annotation.
 80                # seg_filename = os.path.join(
 81                #     FLAGS.semantic_segmentation_folder,
 82                #     filenames[i] + '.' + FLAGS.label_format)
 83                seg_data = tf.gfile.FastGFile(all_label_files[i], 'rb').read()
 84                seg_height, seg_width = label_reader.read_image_dims(seg_data)
 85                if height != seg_height or width != seg_width:
 86                    raise RuntimeError(
 87                        'Shape mismatched between image and label.')
 88                # Convert to tf example.
 89                file_name = os.path.basename(all_image_files[i]).split('.')[0]
 90                example = build_data.image_seg_to_tfexample(
 91                    image_data, file_name, height, width, seg_data)
 92                tfrecord_writer.write(example.SerializeToString())
 93        sys.stdout.write('\n')
 94        sys.stdout.flush()
 95
 96
 97def main(unused_argv):
 98    d = './'
 99    _convert_dataset(d)
100
101if __name__ == '__main__':
102    tf.app.run()

这个代码十分重要，直接将images和masks生成了相应的tfrecord。这个代码中有一些关于路径处理需要根据大家的实际需求修改，再次不敖诉。

deeplab搭建

clone一下models的代码，然后把slim添加到PYTHON 里面：

1export PYTHONPATH=$PYTHONPATH:/home/models/research/slim

在models/research下执行：

1python3 deeplab/model_test.py

没有问题的话就ok了。接下来要修改两个文件：

segmentation_dataset.py
train_utils.py

修改segmentation_dataset.py的110行,添加对应的关于_CAMVID数据集的描述设置：

CamVid描述配置

1# segmentation_dataset.py line 110
2_CAMVID_INFORMATION = DatasetDescriptor(
3    splits_to_sizes={
4        'train': 367,  # num of samples in images/training
5        'val': 101,  # num of samples in images/validation
6    },
7    num_classes=12,
8    ignore_label=255,
9)

这里我以CamVid数据集作为演示。这里的类别是你的分割的类别，但是得加上ignore_label.
CamVid的列别是11类，加上ignore就是12类。

另外在加上：

1DATASETS_INFORMATION = {
2    'cityscapes': _CITYSCAPES_INFORMATION,
3    'pascal_voc_seg': _PASCAL_VOC_SEG_INFORMATION,
4    'ade20k': _ADE20K_INFORMATION,
5    'mydata':_MYDATA_INFORMATION, #我自己的数据集
6    'camvid':_CAMVID_INFORMATION, #camvid示例
7}

注册一下自己的数据集。更改就这两处，大家代开文件之后会看到。
更改一下train_utils.py:

1exclude_list = ['global_step','logits']
2if not initialize_last_layer:
3    exclude_list.extend(last_layers)

目的是在加载与训练的模型的时候不需要top，这个最好更改一下，因为你的类别跟cityscapes不同。

Train

万事具备只欠东风了。直接开始训练，训练的执行路径在 models/research下面，我们运行：

 1python3 deeplab/train.py \
 2    --logtostderr \
 3    --training_number_of_steps=300 \
 4    --train_split='train' \
 5    --model_variant='xception_65' \
 6    --atrous_rates=6 \
 7    --atrous_rates=12 \
 8    --atrous_rates=18 \
 9    --output_stride=16 \
10    --decoder_output_stride=4 \
11    --train_crop_size=513 \
12    --train_crop_size=513 \
13    --train_batch_size=2 \
14    --dataset='minieye_lane' \
15    --tf_initial_checkpoint='deeplab/checkpoints/deeplabv3_cityscapes_train/model.ckpt'\
16    --train_logdir='./deeplab/log' \
17    --dataset_dir='./deeplab/tfrecord'