YOLOv3 代码详解(7) —— 读取pb文件的测试

前言：

yolo系列的论文阅读
论文阅读 || 深度学习之目标检测 重磅出击YOLOv3
论文阅读 || 深度学习之目标检测yolov2
论文阅读 || 深度学习之目标检测yolov1
 
该篇讲解的工程连接是：
tensorflow的yolov3：https://github.com/YunYang1994/tensorflow-yolov3
 
自己对该工程的解析博客：
YOLOv3 || 1. 使用自己的数据集训练yolov3
YOLOv3 || 2. dataset.py解析
YOLOv3 || 3. dataset.py的多进程改造
YOLOv3 || 4. yolov3.py 网络结构的搭建和loss的定义
YOLOv3 || 5. train.py
YOLOv3 || 6. anchorboxes的获取 kmeans.py
YOLOv3 || 7. yolov3的pb文件的测试
 
该篇博客介绍的是yolov3的模型的freeze成pb文件、实际测试调用pb模型

1 代码详解

1.1 freeze_graph.py：将chpt转化为pb文件

需要了解的接口（graph的序列话，及下面的接口的说明在自己的另外一片博客中graph的相关操作）：

• tf.graph_util.convert_variables_to_constants()
  该函数，会将计算图中的变量取值以常量的形式保存。在保存模型文件的时候，这里保存了从输入层到输出层的计算过程的Graph_Def。其余所有不涉及的前向传播和所有的反向传播（梯度）都会被舍弃。
  在保存的时候，参数output_node_names指定保存的节点名称（不是张量的名称）
• sess.graph.as_graph_def()
  将图进行序列化
• graph_def.SerializeToString()
  将序列化的模型转换为字符串，用于写入文件。

import tensorflow as tf
from core.yolov3 import YOLOV3

pb_file = "./yolov3_coco.pb"
ckpt_file = "./checkpoint/yolov3_coco_demo.ckpt"
output_node_names = ["input/input_data", "pred_sbbox/concat_2", "pred_mbbox/concat_2", "pred_lbbox/concat_2"]

# 构建yolov3的网络结构图
with tf.name_scope('input'):
   input_data = tf.placeholder(dtype=tf.float32, name='input_data')
model = YOLOV3(input_data, trainable=False)

# 创建会话，并从ckpt文件中加载权重
sess  = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
saver = tf.train.Saver()
saver.restore(sess, ckpt_file)

# 将graph序列话，并只保留由输入到输出的计算的前向通道
# 注意output_node_names为op的列表
converted_graph_def = tf.graph_util.convert_variables_to_constants(sess,
                           input_graph_def  = sess.graph.as_graph_def(),
                           output_node_names = output_node_names)

with tf.gfile.GFile(pb_file, "wb") as f:
   f.write(converted_graph_def.SerializeToString())

1.2 images_demo.py：读取pb文件并进行预测

import cv2
import numpy as np
import core.utils as utils
import tensorflow as tf
from PIL import Image

return_elements = ["input/input_data:0", "pred_sbbox/concat_2:0", "pred_mbbox/concat_2:0", "pred_lbbox/concat_2:0"]
pb_file         = "./yolov3_coco.pb"
image_path      = "./docs/images/road.jpeg"
num_classes     = 80
input_size      = 416
graph           = tf.Graph()

original_image = cv2.imread(image_path)
original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)
original_image_size = original_image.shape[:2]

# 数据预处理，具体函数后面解析
image_data = utils.image_preporcess(np.copy(original_image), [input_size, input_size])
image_data = image_data[np.newaxis, ...]

# 读取pb文件，具体函数后面解析
return_tensors = utils.read_pb_return_tensors(graph, pb_file, return_elements)


with tf.Session(graph=graph) as sess:
   pred_sbbox, pred_mbbox, pred_lbbox = sess.run(
       [return_tensors[1], return_tensors[2], return_tensors[3]],
               feed_dict={ return_tensors[0]: image_data})

pred_bbox = np.concatenate([np.reshape(pred_sbbox, (-1, 5 + num_classes)),
                           np.reshape(pred_mbbox, (-1, 5 + num_classes)),
                           np.reshape(pred_lbbox, (-1, 5 + num_classes))], axis=0)

bboxes = utils.postprocess_boxes(pred_bbox, original_image_size, input_size, 0.3)

# 使用非极大值抑制进行后处理，具体函数后面解析
bboxes = utils.nms(bboxes, 0.45, method='nms')
image = utils.draw_bbox(original_image, bboxes)
image = Image.fromarray(image)
image.show()

1.2.1 def read_pb_return_tensors() 读取pb文件

导入pb文件思路：

• 使用tf.Graph、tf.GraphDef分别创建个Graph、GraphDef
• 用tensorflow的api tf.gfile.FastGFile(pb_file, 'rb')打开pb文件，并使用.ParseFromString(f.read())导入到GraphDef中
• 使用tf.import_graph_def将GraphDef中的网络结构和参数导入到Graph中。

# pb_file         = "./yolov3_coco.pb"
# return_elements = ["input/input_data:0", "pred_sbbox/concat_2:0", >"pred_mbbox/concat_2:0", "pred_lbbox/concat_2:0"]
# graph           = tf.Graph()
def read_pb_return_tensors(graph, pb_file, return_elements):

   with tf.gfile.FastGFile(pb_file, 'rb') as f:
       frozen_graph_def = tf.GraphDef()
       frozen_graph_def.ParseFromString(f.read())

   with graph.as_default():
       return_elements = tf.import_graph_def(frozen_graph_def,
                                             return_elements=return_elements)
   return return_elements

1.2.2 def image_preporcess() 数据预处理

该函数在yolov3 的数据处理中已经使用过（缩放和填充图片到target_shape）。
 

已知神经网络设定的输入数据的大小为target_shape。 def parse_annotation()：将图片缩放并填充到target_size，并以相同缩放或填充规则处理bboxes，此时bboxes为【左上角-右下角】的形式。
 
具体缩放填充的方式为：

• 缩放：计算图片的长边与target_shape的比值，然后用该比值对原图进行缩放（这样会保持原图的长宽原有比例）
• 填充：然后将图片的短边填充到target_shape

   def image_preporcess(image, target_size, gt_boxes=None):

       image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)

       ih, iw    = target_size
       h,  w, _  = image.shape

       scale = min(iw/w, ih/h)   # 获取长边与target_shape的比值
       nw, nh  = int(scale * w), int(scale * h)  # 获取缩放后的图片的shape
       image_resized = cv2.resize(image, (nw, nh)) # 缩放图片

       image_paded = np.full(shape=[ih, iw, 3], fill_value=128.0) # 创建个target_shape大小的图片，并用128填充
       dw, dh = (iw - nw) // 2, (ih-nh) // 2 
       image_paded[dh:nh+dh, dw:nw+dw, :] = image_resized # 将缩放后的图片放在image_paded的中间位置
       image_paded = image_paded / 255.   # 将图片除以255，进行归一化

       if gt_boxes is None:
           return image_paded

       else:  # 将缩放和填充的操作，应用到bboxes上
           gt_boxes[:, [0, 2]] = gt_boxes[:, [0, 2]] * scale + dw
           gt_boxes[:, [1, 3]] = gt_boxes[:, [1, 3]] * scale + dh
           return image_paded, gt_boxes

1.2.3 def postprocess_boxes() 预测框的后处理

def postprocess_boxes(pred_bbox, org_img_shape, input_size, score_threshold):
   print(pred_bbox.shape)

   valid_scale=[0, np.inf]
   pred_bbox = np.array(pred_bbox)

   pred_xywh = pred_bbox[:, 0:4]
   pred_conf = pred_bbox[:, 4]
   pred_prob = pred_bbox[:, 5:]

   ## (1) 将坐标由(x,y,w,h)，转换到【左上角右下角】的坐标形式
   pred_coor = np.concatenate([pred_xywh[:, :2] - pred_xywh[:, 2:] * 0.5,
                               pred_xywh[:, :2] + pred_xywh[:, 2:] * 0.5], axis=-1)

   ## (2) 将boxes在输入尺寸的坐标，转化到原图尺寸上的坐标
   org_h, org_w = org_img_shape
   resize_ratio = min(input_size / org_w, input_size / org_h)

   dw = (input_size - resize_ratio * org_w) / 2   # 获取到填充边的宽度
   dh = (input_size - resize_ratio * org_h) / 2   # 获取到填充边的宽度

   pred_coor[:, 0::2] = 1.0 * (pred_coor[:, 0::2] - dw) / resize_ratio
   pred_coor[:, 1::2] = 1.0 * (pred_coor[:, 1::2] - dh) / resize_ratio

   ## (3) 去除错误框
   pred_coor = np.concatenate([np.maximum(pred_coor[:, :2], [0, 0]),
                               np.minimum(pred_coor[:, 2:], [org_w - 1, org_h - 1])], axis=-1)
   invalid_mask = np.logical_or((pred_coor[:, 0] > pred_coor[:, 2]), (pred_coor[:, 1] ,pred_coor[:, 3]))
   pred_coor[invalid_mask] = 0
   
   ## (4) 保留面积为正的框
   bboxes_scale = np.sqrt(np.multiply.reduce(pred_coor[:, 2:4] - pred_coor[:, 0:2], axis=-1))
   scale_mask = np.logical_and((valid_scale[0] < bboxes_scale), (bboxes_scale < valid_scale[1]))

   ## (5) 去除得分很低的框
   classes = np.argmax(pred_prob, axis=-1)
   scores = pred_conf * pred_prob[np.arange(len(pred_coor)), classes]
   score_mask = scores > score_threshold

   ## 获取最终保留框的mask，以此获取最后的框的相应的信息
   mask = np.logical_and(scale_mask, score_mask)
   coors, scores, classes = pred_coor[mask], scores[mask], classes[mask]

   return np.concatenate([coors, scores[:, np.newaxis], classes[:, np.newaxis]], axis=-1)

1.2.4 def nms() 非极大值抑制

在yolo系列的论文阅读中，有介绍到非极大值抑制
 

具体操作步骤：(一张图片上有猫，假设有N个框，每个框被分类器计算得分Si，1<=i<=N)
a) 构建一个集合H，初始化包括N个框；构建一个集合O，初始化空集
b) 将集合H中的框按得分大小排序，将得分最高的框m，从集合H移至O
c) 遍历集合H中的框，分别与框m计算IoU。如果高于阈值（一般为0.3~0.5），则认为与此框重叠，将此框从集合H删除
d) 回到第1步，进行迭代，直至集合H为空。集合O中的框为我们所需（集合O中有几个框，图片中就有几只猫）

def nms(bboxes, iou_threshold, sigma=0.3, method='nms'):
   """
   :param bboxes: (xmin, ymin, xmax, ymax, score, class)

   Note: soft-nms, https://arxiv.org/pdf/1704.04503.pdf
         https://github.com/bhara a bvc tsingh430/soft-nms
   """
   # 获取图片中目标物体共有多少种
   classes_in_img = list(set(bboxes[:, 5]))
   best_bboxes = []

   # 针对相同种类的bbox进行计算nms
   for cls in classes_in_img:
       cls_mask = (bboxes[:, 5] == cls)
       cls_bboxes = bboxes[cls_mask]

       while len(cls_bboxes) > 0:

           # 找到得分最高的box
           max_ind = np.argmax(cls_bboxes[:, 4])
           best_bbox = cls_bboxes[max_ind]
           best_bboxes.append(best_bbox)

           # 从cls_bboxes中去除得分最高的框
           cls_bboxes = np.concatenate([cls_bboxes[: max_ind], cls_bboxes[max_ind + 1:]])
           # 计算得分最高的框与其余框的iou
           iou = bboxes_iou(best_bbox[np.newaxis, :4], cls_bboxes[:, :4])

           # 保存着标志符
           weight = np.ones((len(iou),), dtype=np.float32)

           assert method in ['nms', 'soft-nms']
           if method == 'nms':
               iou_mask = iou > iou_threshold
               weight[iou_mask] = 0.0    # iou大于阈值的标识符为0
           if method == 'soft-nms':
               weight = np.exp(-(1.0 * iou ** 2 / sigma))

           # 根据标识符，将大于阈值的boxes的得分置零
           cls_bboxes[:, 4] = cls_bboxes[:, 4] * weight

           score_mask = cls_bboxes[:, 4] > 0.
           # 根据得分，将boxes的信息置零或保存
           cls_bboxes = cls_bboxes[score_mask]

   return best_bboxes