[[206688]]

導(dǎo)語(yǔ)

Tensorflow在更新1.0版本之后多了很多新功能，其中放出了很多用tf框架寫(xiě)的深度網(wǎng)絡(luò)結(jié)構(gòu)(https://github.com/tensorflow/models )，大大降低了開(kāi)發(fā)難度，利用現(xiàn)成的網(wǎng)絡(luò)結(jié)構(gòu)，無(wú)論fine-tuning還是重新訓(xùn)練方便了不少。最近筆者終于跑通TensorFlow Object Detection API的ssd_mobilenet_v1模型，這里記錄下如何完整跑通數(shù)據(jù)準(zhǔn)備到模型使用的整個(gè)過(guò)程，相信對(duì)自己和一些同學(xué)能有所幫助。

Object Detection API提供了5種網(wǎng)絡(luò)結(jié)構(gòu)的預(yù)訓(xùn)練的權(quán)重，全部是用COCO數(shù)據(jù)集進(jìn)行訓(xùn)練，這五種模型分別是SSD+mobilenet、SSD+inception_v2、R-FCN+resnet101、faster RCNN+resnet101、faster RCNN+inception+resnet101。各個(gè)模型的精度和計(jì)算所需時(shí)間如下。下面及介紹下如何使用Object Detection去訓(xùn)練自己的模型。

這里TensorFlow的安裝就不再說(shuō)明了，網(wǎng)上的教程一大把，大家可以找到很詳盡的安裝TensorFlow的文檔。

訓(xùn)練前準(zhǔn)備：

使用protobuf來(lái)配置模型和訓(xùn)練參數(shù)，所以API正常使用必須先編譯protobuf庫(kù)，這里可以下載直接編譯好的pb庫(kù)(https://github.com/google/protobuf/releases )，解壓壓縮包后，把protoc加入到環(huán)境變量中：

$ cd tensorflow/models 
 
$ protoc object_detection/protos/*.proto --python_out=.  

(我是把protoc加到環(huán)境變量中，遇到找不到*.proto文件的報(bào)錯(cuò)，后來(lái)把protoc.exe放到models/object_detection目錄下，重新執(zhí)行才可以)

然后將models和slim(tf高級(jí)框架)加入python環(huán)境變量：

PYTHONPATH=$PYTHONPATH:/your/path/to/tensorflow/models:/your/path/to/tensorflow/models/slim 

數(shù)據(jù)準(zhǔn)備：

數(shù)據(jù)集需要轉(zhuǎn)化成PASCAL VOC結(jié)構(gòu)，API提供了create_pascal_tf_record.py，把VOC結(jié)構(gòu)數(shù)據(jù)集轉(zhuǎn)換成.record格式。不過(guò)我們發(fā)現(xiàn)更簡(jiǎn)單的方式，Datitran提供一種更簡(jiǎn)單生產(chǎn).record格式的方法。

首先需要先要標(biāo)注圖像相應(yīng)標(biāo)簽，這里可以使用labelImg工具。每標(biāo)注一張樣本，即生成一個(gè)xml的標(biāo)注文件。然后，把這些標(biāo)注的xml文件，按訓(xùn)練集與驗(yàn)證集分別放置到兩個(gè)目錄下，在Datitran提供了xml_to_csv.py腳本。這里只要指定標(biāo)注的目錄名即可。接下來(lái)，然后需要我們把對(duì)應(yīng)的csv格式轉(zhuǎn)換成.record格式。

def main(): 
    # image_path = os.path.join(os.getcwd(), 'annotations') 
    image_path = r'D:\training-sets\object-detection\sunglasses\label\test' 
    xml_df = xml_to_csv(image_path) 
    xml_df.to_csv('sunglasses_test_labels.csv', index=None) 
    print('Successfully converted xml to csv.')  

調(diào)用generate_tfrecord.py，注意要指定–csv_input與–output_path這兩個(gè)參數(shù)。執(zhí)行下面命令：

python generate_tfrecord.py --csv_input=sunglasses_test_labels.csv --output_path=sunglass_test.record 

這樣就生成了訓(xùn)練及驗(yàn)證用的train.record與test.record。接下來(lái)指定標(biāo)簽名稱，仿照models/ object_detection/data/ pet_label_map.pbtxt，重新創(chuàng)建一個(gè)文件，指定標(biāo)簽名。

item { 
  id: 1 
  name: 'sunglasses' 
}  

訓(xùn)練：

根據(jù)自己的需要，選擇一款用coco數(shù)據(jù)集預(yù)訓(xùn)練的模型，把前綴model.ckpt放置在待訓(xùn)練的目錄，這里meta文件保存了graph和metadata，ckpt保存了網(wǎng)絡(luò)的weights，這幾個(gè)文件表示預(yù)訓(xùn)練模型的初始狀態(tài)。

打開(kāi)ssd_mobilenet_v1_pets.config文件，并做如下修改：

num_classes:修改為自己的classes num

將所有PATH_TO_BE_CONFIGURED的地方修改為自己之前設(shè)置的路徑(共5處)

其他參數(shù)均保持默認(rèn)參數(shù)。

準(zhǔn)備好上述文件后就可以直接調(diào)用train文件進(jìn)行訓(xùn)練。

python object_detection/train.py \ 
--logtostderr \ 
--pipeline_config_path= D:/training-sets /data-translate/training/ssd_mobilenet_v1_pets.config \ 
--train_dir=D:/training-sets/data-translate/training  

TensorBoard監(jiān)控：

通過(guò)tensorboard工具，可以監(jiān)控訓(xùn)練過(guò)程，輸入西面指令后，在瀏覽器輸入localhost:6006(默認(rèn))即可。

tensorboard --logdir= D:/training-sets/data-translate/training  

這里面有很多指標(biāo)曲線，甚至有模型網(wǎng)絡(luò)架構(gòu)，筆者對(duì)于這里面很多指標(biāo)含義還沒(méi)有弄明白，不過(guò)感覺(jué)出TensorBoard這個(gè)工具應(yīng)該是極其強(qiáng)大。不過(guò)我們可以通過(guò)Total_Loss來(lái)看整體訓(xùn)練的情況。

從整體上看，loss曲線確實(shí)是收斂的，整體的訓(xùn)練效果還是滿意的。另外，TensorFlow還提供了訓(xùn)練過(guò)程中利用驗(yàn)證集驗(yàn)證準(zhǔn)確性的能力，但是筆者在調(diào)用時(shí)，仍有些問(wèn)題，這里暫時(shí)就不詳細(xì)說(shuō)明了。

Freeze Model模型導(dǎo)出：

查看模型實(shí)際的效果前，我們需要把訓(xùn)練的過(guò)程文件導(dǎo)出，生產(chǎn).pb的模型文件。本來(lái)，tensorflow/python/tools/freeze_graph.py提供了freeze model的api，但是需要提供輸出的final node names(一般是softmax之類的***一層的激活函數(shù)命名)，而object detection api提供提供了預(yù)訓(xùn)練好的網(wǎng)絡(luò)，final node name并不好找，所以object_detection目錄下還提供了export_inference_graph.py。

python export_inference_graph.py \ 
--input_type image_tensor 
--pipeline_config_path D:/training-sets /data-translate/training/ssd_mobilenet_v1_pets.config \ 
--trained_checkpoint_prefix D:/training-sets /data-translate/training/ssd_mobilenet_v1_pets.config /model.ckpt-* \ 
--output_directory D:/training-sets /data-translate/training/result  

導(dǎo)出完成后，在output_directory下，會(huì)生成frozen_inference_graph.pb、model.ckpt.data-00000-of-00001、model.ckpt.meta、model.ckpt.data文件。

調(diào)用生成模型：

目錄下本身有一個(gè)調(diào)用的例子，稍微改造如下：

import cv2 
import numpy as np 
import tensorflow as tf 
from object_detection.utils import label_map_util 
from object_detection.utils import visualization_utils as vis_util 
 
 
class TOD(object): 
    def __init__(self): 
        self.PATH_TO_CKPT = r'D:\lib\tf-model\models-master\object_detection\training\frozen_inference_graph.pb' 
        self.PATH_TO_LABELS = r'D:\lib\tf-model\models-master\object_detection\training\sunglasses_label_map.pbtxt' 
        self.NUM_CLASSES = 1 
        self.detection_graph = self._load_model() 
        self.category_index = self._load_label_map() 
 
    def _load_model(self): 
        detection_graph = tf.Graph() 
        with detection_graph.as_default(): 
            od_graph_def = tf.GraphDef() 
            with tf.gfile.GFile(self.PATH_TO_CKPT, 'rb') as fid: 
                serialized_graph = fid.read() 
                od_graph_def.ParseFromString(serialized_graph) 
                tf.import_graph_def(od_graph_def, name='') 
        return detection_graph 
 
    def _load_label_map(self): 
        label_map = label_map_util.load_labelmap(self.PATH_TO_LABELS) 
        categories = label_map_util.convert_label_map_to_categories(label_map, 
                                                                    max_num_classes=self.NUM_CLASSES, 
                                                                    use_display_name=True) 
        category_index = label_map_util.create_category_index(categories) 
        return category_index 
 
    def detect(self, image): 
        with self.detection_graph.as_default(): 
            with tf.Session(graph=self.detection_graph) as sess: 
                # Expand dimensions since the model expects images to have shape: [1, None, None, 3] 
                image_np_expanded = np.expand_dims(image, axis=0) 
                image_tensor = self.detection_graph.get_tensor_by_name('image_tensor:0') 
                boxes = self.detection_graph.get_tensor_by_name('detection_boxes:0') 
                scores = self.detection_graph.get_tensor_by_name('detection_scores:0') 
                classes = self.detection_graph.get_tensor_by_name('detection_classes:0') 
                num_detections = self.detection_graph.get_tensor_by_name('num_detections:0') 
                # Actual detection. 
                (boxes, scores, classes, num_detections) = sess.run( 
                    [boxes, scores, classes, num_detections], 
                    feed_dict={image_tensor: image_np_expanded}) 
                # Visualization of the results of a detection. 
                vis_util.visualize_boxes_and_labels_on_image_array( 
                    image, 
                    np.squeeze(boxes), 
                    np.squeeze(classes).astype(np.int32), 
                    np.squeeze(scores), 
                    self.category_index, 
                    use_normalized_coordinates=True, 
                    line_thickness=8) 
 
        cv2.namedWindow("detection", cv2.WINDOW_NORMAL) 
        cv2.imshow("detection", image) 
        cv2.waitKey(0) 
 
if __name__ == '__main__': 
    image = cv2.imread('image.jpg') 
    detecotr = TOD() 
    detecotr.detect(image)  

下面是一些圖片的識(shí)別效果：

End.