TensorFlow使用object detection训练自己的模型用于物体识别
使用object detection训练并识别自己的模型
1.安装tensorflow(version>=1.4.0)
2.部署tensorflow models
- 在这里下载
- 解压并安装
- 解压后重命名为models复制到tensorflow/目录下
- 在linux下
- 进入tensorflow/models/research/目录,运行protoc object_detection/protos/*.proto --python_out=.
- 在~/.bashrc file.中添加slim和models/research路径
export PYTHONPATH=$PYTHONPATH:/path/to/slim:/path/to/research
- 在windows下
- 下载protoc-3.3.0-win32.zip(version==3.3,已知3.5版本会报错)
- 解压后将protoc.exe放入C:\Windows下
- 在tensorflow/models/research/打开powershell,运行protoc object_detection/protos/*.proto --python_out=.
3.训练数据准备(标记分类的图片)
- 安装labelImg 用来手动标注图片 ,图片需要是png或者jpg格式
- 标注信息会被保存为xml文件,使用 这个脚本 将所有xml文件转换为一个csv文件(xml文件路径识别在29行,根据情况自己修改)
- 把生成的csv文件分成训练集和测试集
4.生成TFRecord文件
- 使用 这个脚本 将两个csv文件生成出两个TFRecord文件(训练自己的模型,必须使用TFRecord格式文件。图片路径识别在86行,根据情况自己修改)
5.创建label map文件
id需要从1开始,class-N便是自己需要识别的物体类别名,文件后缀为.pbtxt
item{
id:1
name: 'class-1'
}
item{
id:2
name: 'class-2'
}
6.下载模型并配置文件
- 下载一个模型(文件后缀.tar.gz)
- 修改对应的训练pipline配置文件
- 查找文件中的PATH_TO_BE_CONFIGURED字段,并做相应修改
- num_classes 改为你模型中包含类别的数量
- fine_tune_checkpoint 解压.tar.gz文件后的路径 + /model.ckpt
- from_detection_checkpoint:true
- train_input_reader
- input_path 由train.csv生成的record格式训练数据
- label_map_path 第5步创建的pbtxt文件路径
- eval_input_reader
- input_path 由test.csv生成的record格式训练数据
- label_map_path 第5步创建的pbtxt文件路径
7. 训练模型
- 进入tensorflow/models/research/目录,运行
python object_detection/train.py --logtostderr --pipeline_config_path=${PATH_TO_YOUR_PIPELINE_CONFIG} //第六步中修改的pipline配置文件路径// --train_dir=${PATH_TO_TRAIN_DIR} //生成的模型保存路径//
8.导出模型
- 在第7步中,--train_dir指向的路径中会生成一系列训练中自动保存的checkpoint,一个checkpoint由三个文件组成,后缀分别是.data-00000-of-00001 .index和.meta,任然在第7步的路径中,运行
python object_detection/export_inference_graph.py --input_type image_tensor --pipeline_config_path ${PIPELINE_CONFIG_PATH} //第六步中修改的pipline配置文件路径// --trained_checkpoint_prefix ${TRAIN_PATH} //上述的一个checkpoint,例如model.ckpt-112254// --output_directory ${OUTPUT_PATH} //输出模型文件的路径//
- 目前直接使用export_infernce_graph.py会报错,解决方法是将object_detection/目录下的export.py文件中第71行的
rewrite_options = rewriter_config_pb2.RewriterConfig(layout_optimizer=rewriter_config_pb2.RewriterConfig.ON)
改为
rewrite_options = rewriter_config_pb2.RewriterConfig()
9.使用新模型识别图片
- 在jupyter notebook中打开object_detection/路径下的object_detection_tutorial.ipynb文件
- 修改第四步代码中的PATH_TO_CKPT,指向第8步导出模型路径中的frozen_inference_graph.pb文件
- 修改第四步代码中的PATH_TO_LABELS,指向第5步创建的pbtxt文件路径
- 修改第四步代码中的NUM_CLASSES为模型实际包含的类别数
- 注释掉第五步所有代码
- 修改第九步代码中的PATH_TO_TEST_IMAGES_DIR, 指向待识别图片文件夹路径
- 修改第九步代码中TEST_IMAGE_PATHS, 将文件名和扩展名修改为实际图片格式,如图片为1.bmp、2.bmp... 则对应改为'{}.bmp',同时修改i值范围以覆盖所有待检测图片
参考文档
https://gist.github.com/douglasrizzo/c70e186678f126f1b9005ca83d8bd2ce
https://towardsdatascience.com/how-to-train-your-own-object-detector-with-tensorflows-object-detector-api-bec72ecfe1d9