因为Android Demo里的模型是已经训练好的,模型保存的label都是固定的,所以我们在使用的时候会发现还有很多东西它识别不出来。那么我们就需要用它来训练我们自己的数据。下面就是使用SSD-MobileNet训练模型的方法。
进入下载的Model目录:
cd models/research/
# 执行
protoc object_detection/protos/*.proto --python_out=.
export PYTHONPATH="${PYTHONPATH}:/home/Github/models:/home/Github/models/research/slim/"
这里的
PYTHONPATH
路径一定要填对,不然会影响到后面运行代码。 然后还要注意代码版本不同文件路径有差别,所以要对照自己目录看好。
继续在research/
目录下执行:
# 如果找不到setup.py, 用find命令找对应的路径
python setup.py build
python setup.py install
在object_detection/
目录下创建目录ssd_model
mkdir object_detection/ssd_model
把下载好的数据集解压进去,数据集路径为
./object_detection/ssd_model/VOCdevkit/
执行配置文件
python ./object_detection/dataset_tools/create_pascal_tf_record.py --label_map_path=object_detection/data/pascal_label_map.pbtxt --data_dir=object_detection/ssd_model/VOCdevkit/ --year=VOC2012 --set=train --output_path=object_detection/ssd_model/pascal_train.record
python ./object_detection/dataset_tools/create_pascal_tf_record.py --label_map_path=object_detection/data/pascal_label_map.pbtxt --data_dir=object_detection/ssd_model/VOCdevkit/ --year=VOC2012 --set=val --output_path=object_detection/ssd_model/pascal_val.record
然后会在ssd_model/
目录下生成pascal_train.record
和pascal_val.record
两个文件,分别有600M左右。
下一步复制训练pet
数据用到的文件,我们在这个基础上修改配置,训练我们的数据
cp object_detection/data/pascal_label_map.pbtxt object_detection/ssd_model/
cp object_detection/samples/configs/ssd_mobilenet_v1_pets.config object_detection/ssd_model/
我们打开pascal_label_map.pbtxt看一下,这个文件里面是类似Json格式的label集,列出了数据集里有哪些label
。Pascal这个数据集label
共有20个。
然后打开配置文件ssd_mobilenet_v1_pets.config,把num_classes
改为20
配置默认训练次数num_steps: 200000
,我们根据自己需要改,注意这个训练是很慢的,差不多以天为单位,所以可以适当改小点。
然后改一些文件路径:
train_input_reader: {
tf_record_input_reader {
input_path: "/home/wow/Github/models/research/object_detection/ssd_model/pascal_train.record"
}
label_map_path: "/home/wow/Github/models/research/object_detection/ssd_model/pascal_label_map.pbtxt"
}
eval_input_reader: {
tf_record_input_reader {
input_path: "/home/wow/Github/models/research/object_detection/ssd_model/pascal_val.record"
}
label_map_path: "/home/wow/Github/models/research/object_detection/ssd_model/pascal_label_map.pbtxt"
shuffle: false
num_readers: 1
}
把之前下载的ssd_mobilenet
解压到/object_detection/ssd_model/ssd_mobilenet
下
把路径填进配置文件
fine_tune_checkpoint: "/home/wow/Github/models/research/object_detection/ssd_model/ssd_mobilenet/model.ckpt"
完成之后,就可以训练模型了
python object_detection/train.py --train_dir object_detection/train --pipeline_config_path object_detection/ssd_model/ssd_mobilenet_v1_pets.config
经过漫长的等待,可以看到在/object_detection/train
目录下生成了模型。然后创建文件夹ssd_model/model
python object_detection/export_inference_graph.py --input_type image_tensor --pipeline_config_path object_detection/ssd_model/ssd_mobilenet_v1_pets.config --trained_checkpoint_prefix object_detection/train/model.ckpt-30000 --output_directory object_detection/ssd_model/model/
生成pb文件,再把pascal_label_map.pbtxt的内容改成.txt
作为label文件,这个模型就可以使用了。
错误1:
TypeError: x and y must have the same dtype, got tf.float32 != tf.int32
修改./object_detection/builders/post_processing_builder.py
def _score_converter_fn_with_logit_scale(tf_score_converter_fn, logit_scale):
"""Create a function to scale logits then apply a Tensorflow function."""
def score_converter_fn(logits):
cr = logit_scale
cr = tf.constant([[cr]],tf.float32)
print(logit_scale)
print(logits)
scaled_logits = tf.divide(logits, cr, name='scale_logits') #change logit_scale
return tf_score_converter_fn(scaled_logits, name='convert_scores')
score_converter_fn.__name__ = '%s_with_logit_scale' % (
tf_score_converter_fn.__name__)
return score_converter_fn
修改之后,需要再执行:
python setup.py build
python setup.py install
错误2:
ImportError: cannot import name rewriter_config_pb2
# 或者
AttributeError: 'module' object has no attribute 'mark_flag_as_required'
修改:
# 前一个错
pip install --upgrade tensorflow==1.2.0
# 后一个错
pip install --upgrade tensorflow==1.4.0
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 = '/home/wow/Github/models/research/object_detection/ssd_model/model/frozen_inference_graph.pb'
self.PATH_TO_LABELS = '/home/wow/Github/models/research/object_detection/ssd_model/pascal_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('/home/wow/Github/models/research/object_detection/ssd_model/img/cow-in-pasture.jpg')
detecotr = TOD()
detecotr.detect(image)
效果:
安装labelImg工具,进行图片的标注.标注后的是xml
格式的文件,将这些文件按照一定比例分到train
和test
目录.
下载datitran,作者自己写了一套xml
转csv
再转为record
文件的脚本.比SSD的脚本方便使用.
首先编辑xml_to_csv.py,修改main函数:
def main():
#image_path = os.path.join(os.getcwd(), 'annotations')
image_path = os.path.join('/home/Github/models/research/object_detection/ssd_model/MyImgs/labels/test')
xml_df = xml_to_csv(image_path)
#xml_df.to_csv('raccoon_labels.csv', index=None)
xml_df.to_csv('fish_test_labels.csv', index=None)
print('Successfully converted xml to csv.')
执行
python xml_to_csv.py
会生成test的csv,同样,修改代码,生成train的csv.
然后进行csv到record的转换
首先修改generate_tfrecord.py
,把main函数的path
改成我们图片路径,然后把if row_label == 'raccoon':
改成我们的label,比如fish
.之后执行下面代码:
python generate_tfrecord.py --csv_input=fish_train_labels.csv --output_path=fish_train.record
python generate_tfrecord.py --csv_input=fish_test_labels.csv --output_path=fish_test.record
然后会生成对应的record文件:
Successfully created the TFRecords: /home/Github/raccoon_dataset/fish_train.record
Successfully created the TFRecords: /home/Github/raccoon_dataset/fish_test.record
回到我们之前训练SSD的目录,创建自己的label文件my_label_map.pbtxt
item {
id: 1
name: 'fish'
}
修改训练配置文件:
num_classes: 1 #20
再把所有PATH_TO_BE_CONFIGURED
的地方改掉,就可以用前面的命令进行训练.
训练时会遇到这个错误:
INFO:tensorflow:Restoring parameters from object_detection/train/model.ckpt-5390
INFO:tensorflow:Error reported to Coordinator: <class 'tensorflow.python.framework.errors_impl.InvalidArgumentError'>, Assign requires shapes of both tensors to match. lhs shape= [1,1,128,12] rhs shape= [1,1,128,126]
这是因为之前我有训练过模型,训练到5390次就停了.这里配置写的是200k次,所以它会接着之前的结果继续跑.但我们的数据发生了变化,所以会出现这个错误.解决方法就是把train目录删掉,重新生成即可