Tensorflow在更新1.0版本之后多了很多新功能,其中放出了很多用tf框架寫的深度網(wǎng)絡結(jié)構,大大降低了開發(fā)難度,利用現(xiàn)成的網(wǎng)絡結(jié)構,無論fine-tuning還是重新訓練方便了不少。而且TensorFlow內(nèi)包含了一個強大的物體檢測API,我們可以利用這API來訓練自己的數(shù)據(jù)集實現(xiàn)特殊的目標檢測。最近筆者終于跑通TensorFlow Object Detection API的ssd_mobilenet_v1模型,這里記錄下如何完整跑通數(shù)據(jù)準備到模型使用的整個過程,相信對大家能有所幫助。
Object Detection API提供了5種網(wǎng)絡結(jié)構的預訓練的權重,全部是用數(shù)據(jù)集進行訓練,這五種模型分別是SSD+mobilenet、SSD+inception_v2、R-FCN+resnet101、faster RCNN+resnet101、faster RCNN+inception+resnet101。各個模型的精度和計算所需時間如下。下面及介紹下如何使用Object Detection去訓練自己的模型。
這里TensorFlow的安裝就不再說明了,網(wǎng)上的教程一大把,大家可以找到很詳盡的安裝TensorFlow的文檔。
tensorflow 訓練模型訓練前準備工作
使用protobuf來配置模型和訓練參數(shù),所以API正常使用必須先編譯protobuf庫,這里可以下載直接編譯好的pb庫( https://github.com/google/protobuf/releases ),解壓壓縮包后,把protoc加入到環(huán)境變量中:
$ cd tensorflow/models
$ protoc object_detection/protos/*.proto --python_out=.
(我是把protoc加到環(huán)境變量中,遇到找不到*.proto文件的報錯,后來把protoc.exe放到models/object_detection目錄下,重新執(zhí)行才可以)
然后將models和slim(tf高級框架)加入python環(huán)境變量:
PYTHONPATH=$PYTHONPATH:/your/path/to/tensorflow/models:/your/path/to/tensorflow/models/slim
tensorflow 訓練模型數(shù)據(jù)準備
數(shù)據(jù)集需要轉(zhuǎn)化成PASCAL VOC結(jié)構,API提供了create_pascal_tf_record.py,把VOC結(jié)構數(shù)據(jù)集轉(zhuǎn)換成.record格式。不過我們發(fā)現(xiàn)更簡單的方式,Datitran提供一種更簡單生產(chǎn).record格式的方法。
首先需要先要標注圖像相應標簽,這里可以使用labelImg工具。每標注一張樣本,即生成一個xml的標注文件。然后,把這些標注的xml文件,按訓練集與驗證集分別放置到兩個目錄下,在Datitran提供了xml_to_csv.py腳本。這里只要指定標注的目錄名即可。接下來,然后需要我們把對應的csv格式轉(zhuǎn)換成.record格式。
def main():
# image_path = os.path.join(os.getcwd(), 'annotations')
image_path = r'D: raining-setsobject-detectionsunglasseslabel est'
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這兩個參數(shù)。執(zhí)行下面命令:
python generate_tfrecord.py --csv_input=sunglasses_test_labels.csv --output_path=sunglass_test.record
這樣就生成了訓練及驗證用的train.record與test.record。接下來指定標簽名稱,仿照models/ object_detection/data/ pet_label_map.pbtxt,重新創(chuàng)建一個文件,指定標簽名。
item {
id: 1
name: 'sunglasses'
}
tensorflow 訓練模型訓練教程
根據(jù)自己的需要,選擇一款用coco數(shù)據(jù)集預訓練的模型,把前綴model.ckpt放置在待訓練的目錄,這里meta文件保存了graph和metadata,ckpt保存了網(wǎng)絡的weights,這幾個文件表示預訓練模型的初始狀態(tài)。
打開ssd_mobilenet_v1_pets.config文件,并做如下修改:
num_classes:修改為自己的classes num
將所有PATH_TO_BE_CONFIGURED的地方修改為自己之前設置的路徑(共5處)
其他參數(shù)均保持默認參數(shù)。
準備好上述文件后就可以直接調(diào)用train文件進行訓練。
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)控:
通過tensorboard工具,可以監(jiān)控訓練過程,輸入西面指令后,在瀏覽器輸入localhost:6006(默認)即可。
tensorboard --logdir= D:/training-sets/data-translate/training
這里面有很多指標曲線,甚至有模型網(wǎng)絡架構,筆者對于這里面很多指標含義還沒有弄明白,不過感覺出TensorBoard這個工具應該是極其強大。不過我們可以通過Total_Loss來看整體訓練的情況。
從整體上看,loss曲線確實是收斂的,整體的訓練效果還是滿意的。另外,TensorFlow還提供了訓練過程中利用驗證集驗證準確性的能力,但是筆者在調(diào)用時,仍有些問題,這里暫時就不詳細說明了。
Freeze Model模型導出:
查看模型實際的效果前,我們需要把訓練的過程文件導出,生產(chǎn).pb的模型文件。本來,tensorflow/python/tools/freeze_graph.py提供了freeze model的api,但是需要提供輸出的final node names(一般是softmax之類的最后一層的激活函數(shù)命名),而object detection api提供提供了預訓練好的網(wǎng)絡,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
導出完成后,在output_directory下,會生成frozen_inference_graph.pb、model.ckpt.data-00000-of-00001、model.ckpt.meta、model.ckpt.data文件。
調(diào)用生成模型:
目錄下本身有一個調(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 f-modelmodels-masterobject_detection rainingfrozen_inference_graph.pb'
self.PATH_TO_LABELS = r'D:lib f-modelmodels-masterobject_detection rainingsunglasses_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,)
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)
下面是一些圖片的識別效果:
評論
查看更多