易 AI - 使用 TensorFlow Object Detection API Mask R-CNN 训练自定义图像分割模型

原文:https://makeoptim.com/en/deep-learning/yiai-image-segmentation

前言

上一篇介绍了目标检测(Object Detection),本文将介绍图像分割(Image Segmentation)的概念,并通过案例讲解如何使用 TensorFlow Object Detection API 来训练自定义的图像分割模型,包括:数据集采集和制作、TensorFlow Object Detection API 安装以及模型的训练

案例效果如下图所示:

image

图像分割

image

如上图所示,目标检测识别图像中存在的内容和检测其位置,而图像分割分为以下三种类型:

  • 语义分割(semantic segmentation):对图像中的每个像素打上类别标签,如上图,把图像分为人(红色)、树木(深绿)、草地(浅绿)、天空(蓝色)标签。
  • 实例分割(instance segmentation)目标检测和语义分割的结合,在图像中将目标检测出来(目标检测),然后对每个像素打上标签(语义分割)。如上图,以人为目标,语义分割不区分属于相同类别的不同实例(所有人都标为红色),实例分割区分同类的不同实例(使用不同颜色区分不同的人)。
  • 全景分割(panoptic segmentation)语义分割和实例分割的结合,即要对所有目标都检测出来,又要区分出同个类别中的不同实例。如上图,实例分割只对图像中的目标(如上图中的人)进行检测和按像素分割,区分不同实例(使用不同颜色),而全景分割是对图中的所有物体包括背景都要进行检测和分割,区分不同实例(使用不同颜色)。

本文要使用的 Mask R-CNN 正是一个实例分割模型。

TensorFlow Object Detection API

TensorFlow Object Detection API 的安装与目标检测一致,参考上一篇 即可,这里不再赘述。

工程创建

注:!!! 从这里开始,请确保在 conda od 的环境下执行。

$ conda activate od
$ conda env list
# conda environments:
#
od                    *  /Users/catchzeng/.conda/envs/od
tensorflow               /Users/catchzeng/.conda/envs/tensorflow
base                     /Users/catchzeng/miniconda3

oda 仓库目录下,执行以下命令,创建工程目录结构。

注:SAVE_DIR 为保存项目的目录,NAME 为项目的名称。

$ make workspace-mask SAVE_DIR=workspace NAME=test-mask
└── workspace
    └── test-mask
        ├── COCO_Image_Viewer.ipynb:用于验证是否成功转换到 COCO format
        ├── Makefile
        ├── annotations:存放标注好的数据集数据(val.record、train.record、label_map.pbtxt)
        ├── create_coco_tf_record.py:将 COCO format 文件转换为 TFRecord
        ├── exported-models:存放训练完之后导出的模型
        ├── exporter_main_v2.py:导出模型脚本
        ├── images:数据集图片和 json 标注
        │   ├── test:手动验证图片
        │   ├── train:训练集图片和 json 标注
        │   └── val:验证集图片和 json 标注
        ├── labelme2coco.py:将 labelme json 文件转换为 COCO format
        ├── model_main_tf2.py:训练模型脚本
        ├── models:自定义模型
        ├── pre-trained-models:[TensorFlow Model Zoo](https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf2_detection_zoo.md) 提供的预训练模型
        └── test_images.py:手动验证图片脚本

数据集

图片

这里还是使用茶杯(cup)、茶壶(teapot)、加湿器(humidifier) 来做案例。

image

将收集的图片,放入工程目录的 images 的三个子目录下。

image

注:本案例只是为了验证如何训练目标识别模型,因此数据集采集得比较少,实际项目中记得尽量采集多点数据集

标注

收集完图片后,需要对训练验证集图片进行标注。标注工具,选用较为常用的 labelme

根据 installation 的说明安装好 labelme,然后执行 labelme 选择 trainval 文件夹进行标注。

image

标注完成后,会生成图片对应的 json 标注文件,如下所示:

workspace/test-mask/images
├── test
│   ├── 15.jpg
│   └── 16.jpg
├── train
│   ├── 1.jpg
│   ├── 1.json
│   ├── 10.jpg
│   ├── 10.json
│   ├── 2.jpg
│   ├── 2.json
│   ├── 3.jpg
│   ├── 3.json
│   ├── 4.jpg
│   ├── 4.json
│   ├── 5.jpg
│   ├── 5.json
│   ├── 6.jpg
│   ├── 6.json
│   ├── 7.jpg
│   ├── 7.json
│   ├── 8.jpg
│   ├── 8.json
│   ├── 9.jpg
│   └── 9.json
└── val
    ├── 11.jpg
    ├── 11.json
    ├── 12.jpg
    ├── 12.json
    ├── 13.jpg
    ├── 13.json
    ├── 14.jpg
    └── 14.json

生成 TFRecord 和 label_map

TensorFlow Object Detection API 只支持 TFRecord 格式的数据集,因此,需要把标注好的数据集进行转换。

cd 到工程目录(cd workspace/test-mask),然后执行 make gen-tfrecord,将在 annotations 文件夹下生成 TFRecord 格式的数据集和 label_map.pbtxt

$ make gen-tfrecord
python labelme2coco.py images/train --output images/train.json
save coco json
images/train.json
python labelme2coco.py images/val --output images/val.json
save coco json
images/val.json
python create_coco_tf_record.py --logtostderr \
        --train_image_dir=images/train \
        --val_image_dir=images/val \
        --train_annotations_file=images/train.json \
        --val_annotations_file=images/val.json \
        --include_masks=True \
        --output_dir=./annotations
......
python gen_label_map.py
item {
  id: 1
  name: 'cup'
}

item {
  id: 2
  name: 'humidifier'
}

item {
  id: 3
  name: 'teapot'
}
annotations
├── label_map.pbtxt
├── train.record
└── val.record

注:如果想验证是否将 labelme json 成功转换到 COCO format,可以运行 COCO_Image_Viewer.ipynb 查看结果。

模型训练

注:!!! 从这里开始,请确保已经 cd 到工程目录(cd workspace/test-mask)。

下载预训练模型

Model Zoo 中选择合适的模型下载解压并放到 workspace/test-mask/pre-trained-models 中。

这里默认选择的是 mask_rcnn_inception_resnet_v2_1024x1024_coco17_gpu-8 可以执行如下命令,自动下载并解压

$ make dl-model

目录结构如下:

└── test-mask
    └── pre-trained-models
        └── mask_rcnn_inception_resnet_v2_1024x1024_coco17_gpu-8
            ├── checkpoint
            ├── pipeline.config
            └── saved_model

配置训练 Pipeline

models 目录创建对应的模型文件夹,比如:mask_rcnn_inception_resnet_v2_1024x1024_coco17_gpu-8,并拷贝 pre-trained-models/mask_rcnn_inception_resnet_v2_1024x1024_coco17_gpu-8/pipeline.config

└── test-mask
    ├── models
    │   └── mask_rcnn_inception_resnet_v2_1024x1024_coco17_gpu-8
    │       └── pipeline.config
    └── pre-trained-models

其中,pipeline.config 如下几处需要根据项目修改

model {
  faster_rcnn {
    num_classes: 3 # 修改为需要识别的目标个数,示例项目为 3 种
    ......
}
train_config {
  batch_size: 8 # 这里需要根据自己的配置,调整大小,这里设置为 8
  num_steps: 10000 # 修改为想要训练的总步数
  optimizer {
    momentum_optimizer {
      learning_rate {
        cosine_decay_learning_rate {
          learning_rate_base: 0.008
          total_steps: 10000 # 修改为想要训练的总步数
          warmup_learning_rate: 0.0
          warmup_steps: 50
        }
      }
      momentum_optimizer_value: 0.8999999761581421
    }
    use_moving_average: false
  }
  gradient_clipping_by_norm: 10.0
  fine_tune_checkpoint_version: V2 # 增加此项,并指定为 V2 版本
  fine_tune_checkpoint: "pre-trained-models/mask_rcnn_inception_resnet_v2_1024x1024_coco17_gpu-8/checkpoint/ckpt-0" # 修改为预制模型的路径
  fine_tune_checkpoint_type: "detection" # 增加此项并设置为 detection,因为我们是做目标检测
  data_augmentation_options {
    random_horizontal_flip {
    }
  }
}
train_input_reader {
  label_map_path: "annotations/label_map.pbtxt" # 增加此项
  tf_record_input_reader {
    input_path: "annotations/train.record" # 修改为训练集的路径
  }
  load_instance_masks: true
  mask_type: PNG_MASKS
}
eval_config {
  metrics_set: "coco_detection_metrics"
  metrics_set: "coco_mask_metrics"
  eval_instance_masks: true
  use_moving_averages: false
  batch_size: 1
  include_metrics_per_category: true
}
eval_input_reader {
  label_map_path: "annotations/label_map.pbtxt" # 修改为标注的路径
  shuffle: false
  num_epochs: 1
  tf_record_input_reader {
    input_path: "annotations/val.record" # 修改为验证集的路径
  }
  load_instance_masks: true
  mask_type: PNG_MASKS
}

训练模型

$ make train

注:如遇以下问题

ValueError: numpy.ndarray size changed, may indicate binary incompatibility. Expected 88 from C header, got 80 from PyObject

执行以下命令,重新安装下 numpy 即可。

pip uninstall numpy
pip install numpy

模型导出

$ make export

测试模型

执行 make export 导出模型后,将测试图片放到 images/test 文件夹下,然后执行 python test_images.py 即可输出标记好目标的图片到 images/test_annotated

image

小结

本文通过案例将图像分割的整个流程都过了一遍,希望能帮助大家快速掌握训练自定义图像分割模型的能力。

案例的代码和数据集都已经放在了 https://github.com/CatchZeng/object-detection-api,有需要的同学可以自行获取。

后面的文章将会为大家带来,目标检测和目标分割的原理以及常用的网络。本篇就到这了,咱们下一篇见。

参考链接

最后编辑于
©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 214,444评论 6 496
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 91,421评论 3 389
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 160,036评论 0 349
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 57,363评论 1 288
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 66,460评论 6 386
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 50,502评论 1 292
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 39,511评论 3 412
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 38,280评论 0 270
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 44,736评论 1 307
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 37,014评论 2 328
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 39,190评论 1 342
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 34,848评论 5 338
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 40,531评论 3 322
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 31,159评论 0 21
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 32,411评论 1 268
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 47,067评论 2 365
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 44,078评论 2 352

推荐阅读更多精彩内容