SSD簡介

SSD（Single Shot MultiBox Detector）是深度學習領域一種新型的目標檢測演算法。在過去的幾次國際比賽中，SSD在速度和準確性方面均取得優異成績，與其他檢測演算法一度拉開很大差距。

SSD的演算法流程大體可以概括為產生候選區域、框選、判斷、過濾幾個步驟。其中，產生候選區域、框選和過濾的演算法是固定的，而針對給定的候選區域，判斷區域中的影象是否是待檢測目標，需要使用不同的模型。常用的識別模型有VGG、Inception、ResNet、MobileNet等。其中，基於MobileNet構建的SSD模型具有最快的檢測速度。

本次訓練使用的是 ssd_mobilenet_v2_coco 這一實現方式，該模型使用新版本的mobile_net，在Coco資料集上進行訓練。

模型的大體工作流程

SSD模型細節非常複雜，但如果只是使用，可以無需瞭解其實現細節。在TensorFlow中，使用SSD模型需要以下幾個步驟：

1. 使用GFile工具類匯入二進位制模型檔案；
2. 將二進位制模型檔案轉換成TensorFlow中的可計算圖（Graph）；
3. 分析圖的結構（或檢視文件）找到圖的資料入口和結果出口；
4. 將資料入口儲存到 tf.Placeholder 類中；
5. 將結果出口儲存到對應的 Tensor 中；
6. 在 Session 中執行 Tensor 對應的 Operation，將真實資料 feed 到對應的 Placeholder 中，然後取出結果資料；
7. 對結果資料進行清洗，視覺化等操作。

模型的訓練方法

標註和準備資料

將待標註的圖片放置在同一個資料夾，隨後下載安裝 labelImg 工具，在工具中開啟待標註資料夾，並選擇結果儲存資料夾，然後開始手動標註圖片。每張圖片對應一個名字相同的XML檔案。

標註完成後，執行 PrepareData.bat，即可生成 TfRecord 檔案。TfRecord 檔案是 TensorFlow 中資料傳輸的標準格式，使用 TfRecord 檔案可以避免小檔案的多次讀取，減少 IO 操作次數。

準備預訓練模型和配置檔案

前往GitHub上下載預訓練模型，然後解壓。

壓縮包中除了預訓練模型，還有和模型配套的配置檔案。配置檔案中大多數內容不需要修改，需要修改的內容如下：

model {
  ssd {
    num_classes: 1
    # 檢測的類別，集裝箱號固定為1類
    image_resizer {
      fixed_shape_resizer {
        height: 300
        width: 300
        # 檢測時圖片大小，預設300*300
      }
    }
  }
}
train_config {
  batch_size: 8
  # 訓練時，每一步的圖片數，記憶體允許的情況下越大越好
  data_augmentation_options {
    random_rgb_to_gray {
      probability : 0.5
    }
  }
  data_augmentation_options {
    random_adjust_brightness {
      max_delta: 0.5
    }
  }
  data_augmentation_options {
    random_adjust_hue {
      max_delta: 0.25
    }
  }
  data_augmentation_options {
    random_adjust_saturation {
    }
  }
  # 上面4個都是對圖片進行隨機變換，提高樣本數
  fine_tune_checkpoint: "object_detection/data/model/model.ckpt"
  # 預訓練模型路徑
  num_steps: 100000
  # 訓練次數
  fine_tune_checkpoint_type: "detection"
}
train_input_reader {
  label_map_path: "object_detection/data/LabelMap.pbtxt"
  # 類別和名稱的對應關係檔案
  tf_record_input_reader {
    input_path: "object_detection/data/mscoco_train.record"
    # 輸入的訓練資料
  }
}

修改完畢後，需要指定類別和名稱的對應關係檔案。檔案內容如下：

item {
  id: 1
  name: 'conNum'
}

隨後即可開始訓練。訓練呼叫 train.py，也可以直接開啟 StartTrain.bat。

匯出模型

訓練過程中會自動儲存模型，但是儲存的是上下文模型，要得到二進位制模型（更精簡，更快速），需要呼叫 export_inference_graph.py。可以點選 ConvertModel.bat 完成操作，也可以直接輸入命令：

python object_detection\export_inference_graph.py \
--input_type image_tensor \
--pipeline_config_path object_detection\data\ModelConfig.config \
--trained_checkpoint_prefix object_detection\data\output\model.ckpt-20000 \
--output_directory object_detection\data\output

模型的使用

程式碼介紹如下：

import tensorflow as tf
from PIL import ImageDraw
from PIL import ImageFont
from PIL import Image
import numpy as np
import time
import os


INPUT_DIR = 'E:\\TrainData\\TestImage'
OUTPUT_DIR = 'E:\\TrainData\\TestResult2'


with tf.gfile.GFile('object_detection/data/output/frozen_inference_graph.pb', 'rb') as file:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(file.read())
with tf.Graph().as_default() as graph:
    tf.import_graph_def(graph_def)
    image_holder = graph.get_tensor_by_name('import/image_tensor:0')
    boxes = graph.get_tensor_by_name('import/detection_boxes:0')
    scores = graph.get_tensor_by_name('import/detection_scores:0')
    num_detections = graph.get_tensor_by_name('import/num_detections:0')
    # classes = graph.get_tensor_by_name('detection_classes:0')
begin = time.time()
with tf.Session(graph=graph) as sess:
    font_en = ImageFont.truetype('C:\\Windows\\Fonts\\batang.ttc', 48)
    for filename in os.listdir(INPUT_DIR):
        image_raw = Image.open(os.path.join(INPUT_DIR, filename))
        image_in = image_raw.resize((300, 300))
        image_in = np.array(image_in, np.uint8)
        if np.shape(image_in)[-1] == 4:
            image_raw = image_raw.convert('RGB')
            image_in = image_in[:, :, :3]
        image_in = np.expand_dims(image_in, axis=0)
        boxes_out, scores_out, nums_out = sess.run([boxes, scores, num_detections], feed_dict={image_holder: image_in})
        height, width, draw = image_raw.height, image_raw.width, ImageDraw.Draw(image_raw)
        for idx in range(int(nums_out)):
            bbox = boxes_out[0][idx]
            draw.rectangle([(bbox[1] * width, bbox[0] * height), (bbox[3] * width, bbox[2] * height)], fill=(0, 0, 0))
            draw.text((bbox[1] * width, bbox[0] * height), str(scores_out[0][idx]), font=font_en, fill=(255, 255, 255))
        if int(nums_out) > 0:
            image_raw.save(os.path.join(OUTPUT_DIR, filename))
            print('Detect: %s' % filename)
        else:
            print('No Result: %s' % filename)
        image_raw.close()
print('Cost: %.2f' % (time.time() - begin))

執行輸出的結果是一個目標區域被畫上黑框的圖片。在訓練一小時以後，檢測率達到70%左右。預計延長訓練時間、加大樣本數量可以大幅度提高準確度。