win10下调用object_detection api ssdmobilenetV2

前言

在Windows平台上去调用object_detection api中模型相比于在Ubuntu下调用更具有难度，在配置的过程中会遇到各种奇奇怪怪的问题。笔者最初尝试用tf2去调用，最终还是放弃改为tf1，总算是成功，特将此经验分享给需要的人。

一、平台环境

1、系统平台：Windows10
2、在anaconda中的虚拟环境中安装python3.6.11，命令为conda create -n yourname pip python=3.6.11；
3、在虚拟环境中安装TensorFlow-gpu 1.15.0，命令为pip install tensorflow-gpu==1.15.0;
4、cuda10.0；
5、cudnn7.6.5；

注意：TensorFlow的版本一定要与cuda、cudnn、python版本对应，不然会出现很多问题。

二、调用步骤

1.下载models文件

上github下载官方model文件，https://github.com/tensorflow/models

2.基础配置

官方文件也会有一直更新，配置方法会有相应的更新，最好的是参考GitHub中官方的方法。
具体配置如下
需要下载protoc文件，具体方法可百度参考其它文章

cd models/research
protoc object_detection/protos/*.proto --python_out=.
cp object_detection/packages/tf1/setup.py .
python -m pip install --use-feature=2020-resolver .

测试是否成功

python object_detection/builders/model_builder_tf1_test.py

如不报错，并且出现OK字样，则说明配置成功

3.训练自己的数据集

1.生成自己的数据文件

可参考其它文章，生成api模型中适用的.tfrecord数据文件

2.模型配置文件修改

我这里使用的是ssd_mobilenet_v2_coco.config模型文件，发现在输入训练命令后一直出现一个奇怪的错误，最后是在Ubuntu系统下进入GitHub复制配置文件到一个新文件中并保存为ssd_mobilenet_v2_coco.config，这才解决了问题。
配置文件中具有修改的有

num_classes: 根据自己的分类数量确定

train_input_reader: {
  tf_record_input_reader {
    input_path: "whsyxt_train.tfrecord"
  }
  label_map_path: "pet_label_map.pbtxt"
}

eval_input_reader: {
  tf_record_input_reader {
    input_path: "whsyxt_validation.tfrecord"
  }
  label_map_path: "pet_label_map.pbtxt"
  shuffle: false
  num_readers: 1
} 

数据路径和验证数据路径根据自己实际路径修改，其它参数也可根据实际情况修改

2.训练

# From the tensorflow/models/research/ directory
PIPELINE_CONFIG_PATH={path to pipeline config file}
MODEL_DIR={path to model directory}
NUM_TRAIN_STEPS=50000
SAMPLE_1_OF_N_EVAL_EXAMPLES=1
python object_detection/model_main.py \
    --pipeline_config_path=${PIPELINE_CONFIG_PATH} \
    --model_dir=${MODEL_DIR} \
    --num_train_steps=${NUM_TRAIN_STEPS} \
    --sample_1_of_n_eval_examples=${SAMPLE_1_OF_N_EVAL_EXAMPLES} \
    --alsologtostderr

3.导出训练模型

python export_inference_graph.py --pipeline_config_path=D:\study\DeepLearning\project\mobilenetSSD\models\research\v2coco.config --trained_checkpoint_prefix ./data/model.ckpt-30000 --output_directory ./data/exported_model_directory

目录根据自己实际情况修改，export_inference_graph.py在models文件中有

4.测试图片

注意在Windows下生成的模型文件名字不要随意修改，笔者修改过后发现会报错，如碰到无法解决的问题，可考虑是否是文件问题，在Ubuntu系统中重新生成。

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'frozen_inference_graph.pb'
        self.PATH_TO_LABELS = r'pet_label_map.pbtxt'
        self.NUM_CLASSES = 20
        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('3.jpg')
    detecotr = TOD()
    detecotr.detect(image)

总结

总的来说，在Windows下会出现更多奇怪的问题，需要考虑好版本的匹配关系，不容易发现的错误是有些文件在Windows下会出现问题，需要多注意。