将训练好的模型转化为pb文件及pb应用

1、相关资料

tensorflow的ckpt文件总结

1.TensorFlow的模型文件

--checkpoint_dir
　　| |--checkpoint
　　| |--MyModel.meta
　　| |--MyModel.data-00000-of-00001
　　| |--MyModel.index

2.meta文件

　　该文件保存的是图结构，meta文件是pb格式，包含变量、结合、OP

3.ckpt文件

　　二进制文件，存储了weights，biases，gradients等变量

4.checkpoint文件

　　文本文件，该文件记录了保存的最新的checkpoint文件以及其他checkpoint文件列表，可以修改这个文件，制定使用哪个model

5.保存TensorFlow

　　使用tf.train.Saver()，TensorFlow中变量都是存储在Session环境中，只有Session环境下才会存有变量值，因此保存模型时需要传入Session。

　　saver=tf.train.Saver()

　　saver.save(sess, './checkpoint_dir/myModel')

6.在实际训练中，我们可能会在每1000次迭代中保存一次模型数据，但是由于图是不变的，没必要每次都去保存，可以通过如下方式指定不保存图

　　saver.save(sess, './checkpoint_dir/MyModel', write_meta_graph=False)

谷歌推荐的保存模型的方式是保存模型为 PB 文件

2、运行脚本及命令

在models/research下执行命令，如果命令太长可以输出为.sh文件

echo  (命令)   >> 名称.sh 

python3 object_detection/export_inference_graph.py --input_type=image_tensor --pipeline_config_path=/root/tf/models/research/object_detection/samples/configs/ssd_resnet50_v1_fpn_shared_box_predictor_640x640_coco14_sync_face.config --trained_checkpoint_prefix=//root/tf/widerface/resnet50v1-fpn/model.ckpt-6214 --output_directory=/root/tf/widerface/resnet50v1-fpn/pb

运行的脚本为/models/research/object_detection/export_inference_graph.py

配置文件--> pipeline_config_path

--pipeline_config_path=/root/tf/models/research/object_detection/samples/configs/ssd_resnet50_v1_fpn_shared_box_predictor_640x640_coco14_sync_face.config

模型--> trained_checkpoint_prefix，后面的数字要保留，显示这是训练了多少次的模型

--trained_checkpoint_prefix=//root/tf/widerface/resnet50v1-fpn/model.ckpt-6214

输出文件夹--> output_directory

--output_directory=/root/tf/widerface/resnet50v1-fpn/pb

3、输出结果

4、应用pb文件

先修改相应pb文件，label_map文件

PATH_TO_FROZEN_GRAPH = "/home/roy/TF/widerface/pb/frozen_inference_graph.pb"

PATH_TOLABELS = "/home/roy/models/research/object_detection/data/face_label_map.pbtxt"

其中frozen_inference_graph.pb为第3步输出结果中的pb文件

脚本文件

import numpy as np
import sys
import tensorflow as tf
import glob
import cv2
sys.path.append("..")


PATH_TO_FROZEN_GRAPH = "/home/roy/TF/widerface/pb/frozen_inference_graph.pb"

PATH_TOLABELS = "/home/roy/models/research/object_detection/data/face_label_map.pbtxt"


# 构图代码
detection_graph = tf.Graph()
with detection_graph.as_default():
        od_graph_def = tf.GraphDef()
        with tf.gfile.GFile(PATH_TO_FROZEN_GRAPH, 'rb') as fid:
          serialized_graph = fid.read()
          od_graph_def.ParseFromString(serialized_graph)
          tf.import_graph_def(od_graph_def, name='')

im_path_list = glob.glob("/home/roy/TF/widerface/test-images/*")  # 获取图片测试的路径
IMAGE_SIZE=(256, 256)
def run_inference_for_single_image(image, graph):
  with graph.as_default():
    with tf.Session() as sess:
      # Get handles to input and output tensors
      ops = tf.get_default_graph().get_operations()
      all_tensor_names = {output.name for op in ops for output in op.outputs}
      tensor_dict = {}
      for key in [
          'num_detections', 'detection_boxes', 'detection_scores',
          'detection_classes', 'detection_masks'
      ]:
        tensor_name = key + ':0'
        if tensor_name in all_tensor_names:
          tensor_dict[key] = tf.get_default_graph().get_tensor_by_name(
              tensor_name)
      if 'detection_masks' in tensor_dict:
        # The following processing is only for single image\n",
        detection_boxes = tf.squeeze(tensor_dict['detection_boxes'], [0])
        detection_masks = tf.squeeze(tensor_dict['detection_masks'], [0])
        # Reframe is required to translate mask from box coordinates to image coordinates and fit the image size.
        real_num_detection = tf.cast(tensor_dict['num_detections'][0], tf.int32)
        detection_boxes = tf.slice(detection_boxes, [0, 0], [real_num_detection, -1])
        detection_masks = tf.slice(detection_masks, [0, 0, 0], [real_num_detection, -1, -1])
        detection_masks_reframed = utils_ops.reframe_box_masks_to_image_masks(
            detection_masks, detection_boxes, image.shape[1], image.shape[2])
        detection_masks_reframed = tf.cast(
            tf.greater(detection_masks_reframed, 0.5), tf.uint8)
        # Follow the convention by adding back the batch dimension
        tensor_dict['detection_masks'] = tf.expand_dims(
            detection_masks_reframed, 0)
      image_tensor = tf.get_default_graph().get_tensor_by_name('image_tensor:0')

      # Run inference
      output_dict = sess.run(tensor_dict,
                             feed_dict={image_tensor: np.expand_dims(image, 0)})
      # all outputs are float32 numpy arrays, so convert types as appropriate
      output_dict['num_detections'] = int(output_dict['num_detections'][0])
      output_dict['detection_classes'] = output_dict[
          'detection_classes'][0].astype(np.int64)
      output_dict['detection_boxes'] = output_dict['detection_boxes'][0]
      output_dict['detection_scores'] = output_dict['detection_scores'][0]
      if 'detection_masks' in output_dict:
        output_dict['detection_masks'] = output_dict['detection_masks'][0]
  return output_dict

for image_path in im_path_list:
    imdata = cv2.imread(image_path)
    sp = imdata.shape
    imdata = cv2.resize(imdata, IMAGE_SIZE)  # 重新定义图片尺寸
    output_dict = run_inference_for_single_image(imdata, detection_graph)
    for i in range(len(output_dict['detection_scores'])):
        if output_dict['detection_scores'][i] > 0.6:  # 人脸框的预值要大于0.6我们才认为它是一个人脸框
            bbox = output_dict['detection_boxes'][i]  # box 人脸框
            y1 = int(IMAGE_SIZE[0] * bbox[0])
            x1 = int(IMAGE_SIZE[1] * bbox[1])
            y2 = int(IMAGE_SIZE[0] * bbox[2])
            x2 = int(IMAGE_SIZE[1] * bbox[3])
            cv2.rectangle(imdata, (x1, y1), (x2, y2), (0, 255, 0), 2) # 绘制人脸框，最后一个值2为线条宽度
    cv2.imshow(image_path, imdata)
    cv2.waitKey(0)
    cv2.destroyWindow(image_path)

结果：