易 AI - 使用 TensorFlow Object Detection API 訓練自定義目標檢測模型

原文：https://makeoptim.com/deep-learning/yiai-object-detection

• 前言
• 目標檢測
  • 位置
• 發展史
  • 傳統方法（候選區域+手工特徵提取+分類器）
  • Region Proposal+CNN（Two-stage）
  • 端到端（One-stage）
• TensorFlow Object Detection API
  • 安裝依賴項
  • 安裝 API
• 工程創建
• 數據集
  • 圖片
  • 標註
  • LabelMap
  • 創建 TFRecord
• 模型訓練
  • 下載預訓練模型
  • 配置訓練 Pipeline
  • 訓練模型
  • 模型導出與轉換
  • 測試模型
• 小結
• 參考鏈接

前言

本文將介紹目標檢測（Object Detection）的概念，並通過案例講解如何使用 TensorFlow Object Detection API 來訓練自定義的目標檢測器，包括：數據集採集和製作、TensorFlow Object Detection API 安裝以及模型的訓練。

案例效果如下圖所示：

目標檢測

如上圖所示，圖像分類解決的問題是圖中的物體是什麼，而目標檢測能識別圖片中有哪些物體以及物體的位置（座標）。

位置

目標檢測的位置信息一般有兩種格式：

• 極座標(xmin, ymin, xmax, ymax)：
  • xmin,ymin：x,y 座標的最小值
  • xmin,ymin：x,y 座標的最大值
• 中心點：(x_center, y_center, w, h)
  • x_center, y_center：目標檢測框的中心點座標
  • w,h：目標檢測框的寬、高

注：圖片左上角爲原點(0,0)

發展史

傳統方法（候選區域+手工特徵提取+分類器）

HOG+SVM、DPM

Region Proposal+CNN（Two-stage）

R-CNN, SPP-NET, Fast R-CNN, Faster R-CNN

端到端（One-stage）

YOLO、SSD

TensorFlow Object Detection API

TensorFlow Object Detection API 是一個構建在 TensorFlow 之上的開源框架，可以輕鬆構建、訓練和部署對象檢測模型。另外，TensorFlow Object Detection API 還提供了 Model Zoo 方便我們選擇和切換預訓練模型。

安裝依賴項

• conda
• protoc

使用以下命令檢查是否安裝成功。

$ conda --version
conda 4.9.2
$ protoc --version
libprotoc 3.17.1

安裝 API

TensorFlow Object Detection API 官方的安裝步驟較爲繁瑣，筆者寫了一個腳本直接一鍵安裝。

執行 git clone https://github.com/CatchZeng/object-detection-api.git 下載倉庫，然後到該倉庫（下文簡稱 oda 倉庫）目錄下，執行以下命令，如果看到如下輸出，表示安裝成功。

$ conda create -n  od python=3.8.5 && conda activate od && make install
......
----------------------------------------------------------------------
Ran 24 tests in 21.869s

OK (skipped=1)

注：如果你不想用 conda，可以在自己的 python 環境上直接使用 make install 安裝即可，比如在 colab 中使用。

注：由於 cudaDNN 和 toolkit 更新可能沒有 TensorFlow 快。因此，如果你的機器有 GPU，安裝完成後，需要將 TensorFlow 降回到 cudaDNN 和 toolkit 支持的版本這樣才能支持 GPU 訓練，以 2.4.1 爲例：

$ pip install --upgrade tf-models-official==2.4.0
$ pip install --upgrade tensorflow==2.4.1

注：如果安裝失敗，可以參考官方文檔的詳細步驟安裝。

工程創建

注：!!! 從這裏開始，請確保在 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-box SAVE_DIR=workspace NAME=test

└── workspace
    └── test
        ├── Makefile
        ├── annotations：存放標註好的數據集數據（val.record、train.record、label_map.pbtxt）
        ├── convert_quant_lite.py：量化 tflite 模型腳本
        ├── export_tflite_graph_tf2.py：導出 tflite 模型腳本
        ├── exported-models：存放訓練完之後導出的模型
        ├── exporter_main_v2.py：導出模型腳本
        ├── images：數據集圖片和 xml 標註
        │   ├── test：手動驗證圖片
        │   ├── train：訓練集圖片和 xml 標註
        │   └── val：驗證集圖片和 xml 標註
        ├── model_main_tf2.py：訓練模型腳本
        ├── models：自定義模型
        ├── pre-trained-models：TensorFlow Model Zoo 提供的預訓練模型
        └── test_images.py：手動驗證圖片腳本

數據集

圖片

筆者喜歡喝茶，這次就用茶杯(cup)、茶壺(teapot)、加溼器(humidifier) 來做案例吧。

將收集的圖片，放入工程目錄的 images 的三個子目錄下。

注：本案例只是爲了驗證如何訓練目標識別模型，因此數據集採集得比較少，實際項目中記得儘量採集多點數據集。

標註

收集完圖片後，需要對訓練和驗證集圖片進行標註。標註工具，選用較爲常用的 LabelImg。

根據 installation 的說明安裝好 LabelImg，然後執行 labelImg 選擇 train 和 val 文件夾進行標註。

標註完成後，會生成圖片對應的 xml 標註文件，如下所示：

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

LabelMap

在 workspace/test/annotations 文件夾下創建 label_map.pbtxt，內容爲模型需要識別的目標。

item {
    id: 1
    name: 'cup'
}

item {
    id: 2
    name: 'teapot'
}

item {
    id: 3
    name: 'humidifier'
}

創建 TFRecord

TensorFlow Object Detection API 只支持 TFRecord 格式的數據集，因此，需要把標註好的數據集進行轉換。

先 cd 到工程目錄（cd workspace/test），然後執行 make gen-tfrecord，將在 annotations 文件夾下生成 TFRecord 格式的數據集。

$ make gen-tfrecord
python ../../scripts/preprocessing/generate_tfrecord.py \
        -x images/train \
        -l annotations/label_map.pbtxt \
        -o annotations/train.record
Successfully created the TFRecord file: annotations/train.record
python ../../scripts/preprocessing/generate_tfrecord.py \
        -x images/val \
        -l annotations/label_map.pbtxt \
        -o annotations/val.record
Successfully created the TFRecord file: annotations/val.record

annotations
├── label_map.pbtxt
├── train.record
└── val.record

模型訓練

注：!!! 從這裏開始，請確保已經 cd 到工程目錄（cd workspace/test）。

下載預訓練模型

從 Model Zoo 中選擇合適的模型下載解壓並放到 workspace/test/pre-trained-models 中。

如果你選擇的是 SSD MobileNet V2 FPNLite 320x320 可以執行如下命令，自動下載並解壓

$ make dl-model

目錄結構如下：

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

配置訓練 Pipeline

在 models 目錄創建對應的模型文件夾，比如：ssd_mobilenet_v2_fpnlite_320x320，並拷貝 pre-trained-models/ssd_mobilenet_v2_fpnlite_320x320_coco17_tpu-8/pipeline.config。

└── test
    ├── models
    │   └── ssd_mobilenet_v2_fpnlite_320x320
    │       └── pipeline.config
    └── pre-trained-models

其中，pipeline.config 如下幾處需要根據項目修改

model {
  ssd {
    num_classes: 3 # 修改爲需要識別的目標個數，示例項目爲 3 種
    ......
}
train_config {
  batch_size: 8 # 這裏需要根據自己的配置，調整大小，這裏設置爲 8
  ......
  optimizer {
    momentum_optimizer {
      learning_rate {
        cosine_decay_learning_rate {
          learning_rate_base: 0.07999999821186066
          total_steps: 10000 # 修改爲想要訓練的總步數
          warmup_learning_rate: 0.026666000485420227
          warmup_steps: 1000
        }
      }
      momentum_optimizer_value: 0.8999999761581421
    }
    use_moving_average: false
  }
  fine_tune_checkpoint: "pre-trained-models/ssd_mobilenet_v2_fpnlite_320x320_coco17_tpu-8/checkpoint/ckpt-0" # 修改爲預製模型的路徑
  num_steps: 10000 # 修改爲想要訓練的總步數
  startup_delay_steps: 0.0
  replicas_to_aggregate: 8
  max_number_of_boxes: 100
  unpad_groundtruth_tensors: false
  fine_tune_checkpoint_type: "detection" # 這裏需要修改爲 detection，因爲我們是做目標檢測
  fine_tune_checkpoint_version: V2
}
train_input_reader {
  label_map_path: "annotations/label_map.pbtxt" # 修改爲標註的路徑
  tf_record_input_reader {
    input_path: "annotations/train.record" # 修改爲訓練集的路徑
  }
}
eval_config {
  metrics_set: "coco_detection_metrics"
  use_moving_averages: false
}
eval_input_reader {
  label_map_path: "annotations/label_map.pbtxt" # 修改爲標註的路徑
  shuffle: false
  num_epochs: 1
  tf_record_input_reader {
    input_path: "annotations/val.record" # 修改爲驗證集的路徑
  }
}

訓練模型

$ 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
  

• TFLite 模型

  $ make export-lite
  

• 轉換 TFLite 模型

  $ make convert-lite
  

• 轉換量化版 TFLite 模型

  $ make convert-quant-lite
  

注：以上命令，大家可以加 -640 表示使用 SSD MobileNet V2 FPNLite 640x640 的模型，比如：make train -> make train-640

測試模型

執行 make export 導出模型後，將測試圖片放到 images/test 文件夾下，然後執行 python test_images.py 即可輸出標記好目標的圖片到 images/test_annotated。

小結

本文通過案例將目標檢測的整個流程都過了一遍，希望能幫助大家快速掌握訓練自定義目標檢測器的能力。

案例的代碼和數據集都已經放在了 https://github.com/CatchZeng/object-detection-api，有需要的同學可以自行獲取。

後面的文章將會爲大家帶來，目標檢測的原理、常用的目標檢測網絡，以及目標分割。本篇就到這了，咱們下一篇見。

參考鏈接

• https://github.com/tensorflow/models/tree/master/research/object_detection
• https://arxiv.org/pdf/1905.05055.pdf
• https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf2_detection_zoo.md
• https://stackoverflow.com/questions/66060487/valueerror-numpy-ndarray-size-changed-may-indicate-binary-incompatibility-exp