前言

之前寫了一篇如何在windows系統上安裝Tensorflow Object Detection API？（點擊跳轉）

然后就想著把數據集換成自己的數據集進行訓練得到自己的目標檢測模型。動手之前先學習了一波別人是如何實現的，看了大多數教程都有一個小問題：用VOC2012數據集進行訓練當做用自己的數據集。

然而，初心想看的是自己的數據集啊！于是就自己來擼一篇教程，方便自己也給別人一些參考吧~

目錄

基于自己數據集進行目標檢測訓練的整體步驟如下：

• 數據標注，制作VOC格式的數據集

• 將數據集制作成tfrecord格式

• 下載預使用的目標檢測模型

• 配置文件和模型

• 模型訓練

這里放一下小詹這個項目的整體截圖，方便后邊文件的對號入座。

數據標注，制作VOC格式的數據集

數據集當然是第一步，在收集好數據后需要進行數據的標注，考慮到VOC風格，這里推薦使用LabelImg工具進行標注。

至于工具具體怎么用，自己摸索下就好，小詹已經把關鍵點地方框選出來啦。（Tip: Ctrl+R選擇標注文件存放路徑）

將數據集制作成tfrecord格式

這一部需要將手動標注的xml文件進行處理，得到標注信息csv文件，之后和圖像數據一起制作成tfrecord格式的數據，用于網絡訓練。

xml轉換為csv文件

這一步需要對xml文件進行解析，提取出標注信息存入csv文件，這里直接把小詹的腳步文件（Xml2Csv.py）分享如下，當然文件路徑你得換成自己的！

#?將xml文件讀取關鍵信息轉化為csv文件

import?os
import?glob
import?pandas?as?pd
import?xml.etree.ElementTree?as?ET


def?xml_to_csv(path):
????xml_list?=?[]
????for?xml_file?in?glob.glob(path?+?'/*.xml'):
????????tree?=?ET.parse(xml_file)
????????root?=?tree.getroot()
????????for?member?in?root.findall('object'):
????????????value?=?(root.find('filename').text,
?????????????????????int(root.find('size')[0].text),
?????????????????????int(root.find('size')[1].text),
?????????????????????member[0].text,
?????????????????????int(member[4][0].text),
?????????????????????int(member[4][1].text),
?????????????????????int(member[4][2].text),
?????????????????????int(member[4][3].text)
?????????????????????)
????????????xml_list.append(value)
????column_name?=?['filename',?'width',?'height',?'class',?'xmin',?'ymin',?'xmax',?'ymax']
????xml_df?=?pd.DataFrame(xml_list,?columns=column_name)
????return?xml_df


def?main():
????image_path?=?r'E:\Jan_Project\Experiment_1\dataset\test_xml'
????xml_df?=?xml_to_csv(image_path)
????xml_df.to_csv(r'E:\Jan_Project\Experiment_1\dataset\cancer_test_labels.csv',?index=None)
????print('Successfully?converted?xml?to?csv.')


if?__name__?==?'__main__':
????main()

生成tfrecord數據文件

之后在對應文件路徑處就有了csv文件，再利用如下腳步自動生成tfrecord。（這是github上生成文件的修改版）

#?將CSV文件和圖像數據整合為TFRecords

"""
name:?generate_tfrecord.py

Usage:
??#?From?tensorflow/models/
??#?Create?train?data:
??python?generate_tfrecord.py?--csv_input=data/train_labels.csv??--output_path=train.record

??#?Create?test?data:
??python?generate_tfrecord.py?--csv_input=data/test_labels.csv??--output_path=test.record
"""
from?__future__?import?division
from?__future__?import?print_function
from?__future__?import?absolute_import

import?os
import?io
import?pandas?as?pd
import?tensorflow?as?tf

from?PIL?import?Image
from?object_detection.utils?import?dataset_util
from?collections?import?namedtuple,?OrderedDict

flags?=?tf.app.flags
flags.DEFINE_string('csv_input',?'',?'Path?to?the?CSV?input')
flags.DEFINE_string('output_path',?'',?'Path?to?output?TFRecord')
FLAGS?=?flags.FLAGS


#?TO-DO?replace?this?with?label?map
def?class_text_to_int(row_label):
????if?row_label?==?'yichang':
????????return?1
????else:
????????None


def?split(df,?group):
????data?=?namedtuple('data',?['filename',?'object'])
????gb?=?df.groupby(group)
????return?[data(filename,?gb.get_group(x))?for?filename,?x?in?zip(gb.groups.keys(),?gb.groups)]


def?create_tf_example(group,?path):
????with?tf.gfile.GFile(os.path.join(path,?'{}'.format(group.filename)),?'rb')?as?fid:
????????encoded_jpg?=?fid.read()
????encoded_jpg_io?=?io.BytesIO(encoded_jpg)
????image?=?Image.open(encoded_jpg_io)
????width,?height?=?image.size

????filename?=?group.filename.encode('utf8')
????image_format?=?b'jpg'
????xmins?=?[]
????xmaxs?=?[]
????ymins?=?[]
????ymaxs?=?[]
????classes_text?=?[]
????classes?=?[]

????for?index,?row?in?group.object.iterrows():
????????xmins.append(row['xmin']?/?width)
????????xmaxs.append(row['xmax']?/?width)
????????ymins.append(row['ymin']?/?height)
????????ymaxs.append(row['ymax']?/?height)
????????classes_text.append(row['class'].encode('utf8'))
????????classes.append(class_text_to_int(row['class']))

????tf_example?=?tf.train.Example(features=tf.train.Features(feature={
????????'image/height':?dataset_util.int64_feature(height),
????????'image/width':?dataset_util.int64_feature(width),
????????'image/filename':?dataset_util.bytes_feature(filename),
????????'image/source_id':?dataset_util.bytes_feature(filename),
????????'image/encoded':?dataset_util.bytes_feature(encoded_jpg),
????????'image/format':?dataset_util.bytes_feature(image_format),
????????'image/object/bbox/xmin':?dataset_util.float_list_feature(xmins),
????????'image/object/bbox/xmax':?dataset_util.float_list_feature(xmaxs),
????????'image/object/bbox/ymin':?dataset_util.float_list_feature(ymins),
????????'image/object/bbox/ymax':?dataset_util.float_list_feature(ymaxs),
????????'image/object/class/text':?dataset_util.bytes_list_feature(classes_text),
????????'image/object/class/label':?dataset_util.int64_list_feature(classes),
????}))
????return?tf_example


def?main(csv_input,?output_path,?image_path):
????writer?=?tf.python_io.TFRecordWriter(output_path)
????path?=?image_path
????examples?=?pd.read_csv(csv_input)
????grouped?=?split(examples,?'filename')
????for?group?in?grouped:
????????tf_example?=?create_tf_example(group,?path)
????????writer.write(tf_example.SerializeToString())

????writer.close()
????print('Successfully?created?the?TFRecords:?{}'.format(output_path))

if?__name__?==?'__main__':
#?????csv_input?=?r'E:\Jan_Project\Experiment_1\dataset\cancer_train_labels.csv'
#?????output_path?=?r'E:\Jan_Project\Experiment_1\dataset\train.tfrecord'
#?????image_path?=?r'E:\Jan_Project\Experiment_1\dataset\train_img'
#?????main(csv_input,?output_path,?image_path)

????csv_input?=?r'E:\Jan_Project\Experiment_1\dataset\cancer_test_labels.csv'
????output_path?=?r'E:\Jan_Project\Experiment_1\dataset\test.tfrecord'
????image_path?=?r'E:\Jan_Project\Experiment_1\dataset\test_img'
????main(csv_input,?output_path,?image_path)

利用上述腳步后便得到了想要的數據格式，小詹這里如圖所示：

下載預使用的目標檢測模型

準備好訓練數據后，選擇模型進行訓練，下載官方預訓練模型【Github】
對于目標檢測，可以考慮選擇幾種最常用的模型：

• ssd_mobilenet_v1_coco

• ssd_mobilenet_v2_coco

• faster_rcnn_resnet50_coco

• faster_rcnn_resnet101_coco

小詹選擇的是上方鏈接中對應下圖的那個，自己視情況而定即可。

下載后解壓到對應文件夾中（見小詹放的第一張項目整體圖）

配置文件和模型

建立label_map.pbtxt

這里需要針對自己數據集進行修改，格式如下：

item{
? ?id: 1
? ?name: 'object'
}

修改

進入tensorflow/models/research/object_detection/samples/config文件夾找到對應自己模型的config文件，針對自己的情況進行修改：

num_classes:?修改為你自己任務的類別數
batch?size：2（GPU顯存較小的，盡量設置成小數值）

fine_tune_checkpoint:?"路徑/model.ckpt"?#指定“訓練模型的檢查點文件”

train_input_reader:?{
??tf_record_input_reader?{
????input_path:?"路徑/train.tfrecord"
??}
??label_map_path:?"路徑/label_map.pbtxt"
}

eval_input_reader:?{
??tf_record_input_reader?{
????input_path:?"路徑/test.tfrecord"
??}
??label_map_path:?"路徑/label_map.pbtxt"
??shuffle:?false
??num_readers:?1

模型訓練

關于訓練，要注意batch size大小和網絡模型復雜程度，注意顯存是否夠大？顯存不夠就OOM（out of memory）了。

訓練模型只需要運行object_detection/legacy路徑下的train.py程序即可。(當然object_detection API安裝是大前提，具體看上一篇文章！)

本地電腦：

python?object_detection//legacy//train.py?--logtostderr?--train_dir=E://Jan_Project//Experiment_1//model?--pipeline_config_path=E://Jan_Project//Experiment_1//training//faster_rcnn_inception_v2_coco.config

如果配置不夠，可以云服務器上跑。以下是訓練過程截圖。

訓練后還可以導出模型,用于檢測測試。

#From?tensorflow/modles/research/object_detection/

python?export_inference_graph.py?
--input_type?image_tensor?
--pipeline_config_path?路徑/***.config
--trained_checkpoint_prefix?路徑/model.ckpt-numbers?#選擇最近的一個或確認收斂到最優的一個
--output_directory?路徑/my_model/?#模型的輸出路徑

以上整理于2019-5-20,節日快樂！

相關推薦：

如何在windows系統上安裝Tensorflow Object Detection API？

總結

以上是生活随笔為你收集整理的如何在Windows系统上使用Object Detection API训练自己的数据？的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。