TenSorRT部署运行yolov3

本文是基于TensorRT 5.0.2基础上，关于其内部的yolov3_onnx例子的分析和介绍。
本例子展示一个完整的ONNX的pipline，在tensorrt 5.0的ONNX-TensorRT基础上，基于Yolov3-608网络进行inference，包含预处理和后处理。

• 首先，从作者网站下载yolov3，然后将其转换成onnx形式，接着基于onnx的graph生成一个tensorrt engine;
• 然后，在样本图片上进行预处理，并将结果作为engine的输入;
• 在inference之后，开始关于包含bounding-box聚类的后处理，然后最终得到一个新的图像文件，并将其存放在磁盘上，以便后续肉眼观察。

1 引言

假设当前路径为：

TensorRT-5.0.2.6/samples

其对应当前例子文件目录树为：

# tree python

python
├── common.py
└── yolov3_onnx
├── coco_labels.txt
├── data_processing.py
├── onnx_to_tensorrt.py
├── README.md
├── requirements.txt
├── yolov3.cfg
├── yolov3.weights
├── dog.jpg
└── yolov3_to_onnx.py

其中：

• yolov3_to_onnx.py：将原始yolov3模型转换成onnx结构。该脚本会自动下载所需要依赖文件；
• onnx_to_tensorrt.py：将onnx的yolov3转换成engine然后进行inference。

2 darknet转onnx

首先运行：

python yolov3_to_onnx.py

就会自动从作者网站下载yolo3的所需依赖

from __future__ import print_function
from collections import OrderedDict
import hashlib
import os.path

import wget

import onnx # github网址为https://github.com/onnx/onnx
from onnx import helper
from onnx import TensorProto
import numpy as np

import sys

‘’‘main第二步：解析yolov3.cfg ‘’’
class DarkNetParser(object):
“”“定义一个基于DarkNet YOLOv3-608的解析器.”""

<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">__init__</span><span class="hljs-params">(self, supported_layers)</span>:</span>
    <span class="hljs-string">"""初始化DarkNetParser对象.

    Keyword argument:
    supported_layers -- 一个list，其中每个元素为字符串，表示支持的层，以DarkNet的命名习惯,
    """</span>

    self.layer_configs = OrderedDict()
    self.supported_layers = supported_layers
    self.layer_counter = <span class="hljs-number">0</span>

<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">parse_cfg_file</span><span class="hljs-params">(self, cfg_file_path)</span>:</span>
    <span class="hljs-string">"""逐层解析yolov3.cfg文件,以字典形式追加每层的参数到layer_configs

    Keyword argument:
    cfg_file_path -- yolov3.cfg文件的路径
    """</span>

    <span class="hljs-keyword">with</span> open(cfg_file_path, <span class="hljs-string">'rb'</span>) <span class="hljs-keyword">as</span> cfg_file:
        remainder = cfg_file.read()
        remainder = remainder.decode(<span class="hljs-string">'utf-8'</span>) <span class="hljs-comment"># 这一行for py3</span>
        <span class="hljs-keyword">while</span> remainder:
            <span class="hljs-comment"># 一次次的去处理字符串，如果返回的layer_dict有值，则表示当前已经处理完一个字典</span>
            layer_dict, layer_name, remainder = self._next_layer(remainder)
            <span class="hljs-keyword">if</span> layer_dict:
                self.layer_configs[layer_name] = layer_dict

    <span class="hljs-keyword">return</span> self.layer_configs

<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">_next_layer</span><span class="hljs-params">(self, remainder)</span>:</span>
    <span class="hljs-string">"""将其视为一个字符串，然后以DarkNet的分隔符来逐段处理.
    在最近的分隔符之后，返回层参数和剩下的字符串
    如文件中第一个Conv层 ...

    [convolutional]
    batch_normalize=1
    filters=32
    size=3
    stride=1
    pad=1
    activation=leaky

    ... 会变成如下形式字典:
    {'activation': 'leaky', 'stride': 1, 'pad': 1, 'filters': 32,
    'batch_normalize': 1, 'type': 'convolutional', 'size': 3}.

    '001_convolutional' 是层名layer_name, 后续所有字符以remainder表示的字符串返回

    Keyword argument:
    remainder -- 仍需要处理的字符串
    """</span>

    <span class="hljs-comment"># head，tail方式</span>
    <span class="hljs-comment"># 读取'['，然后获取tail</span>
    remainder = remainder.split(<span class="hljs-string">'['</span>, <span class="hljs-number">1</span>)
    <span class="hljs-keyword">if</span> len(remainder) == <span class="hljs-number">2</span>:
        remainder = remainder[<span class="hljs-number">1</span>]
    <span class="hljs-keyword">else</span>:
        <span class="hljs-keyword">return</span> <span class="hljs-keyword">None</span>, <span class="hljs-keyword">None</span>, <span class="hljs-keyword">None</span>
    <span class="hljs-comment"># 读取‘]’，然后获取tail</span>
    remainder = remainder.split(<span class="hljs-string">']'</span>, <span class="hljs-number">1</span>)
    <span class="hljs-keyword">if</span> len(remainder) == <span class="hljs-number">2</span>:
        layer_type, remainder = remainder
    <span class="hljs-keyword">else</span>:
        <span class="hljs-keyword">return</span> <span class="hljs-keyword">None</span>, <span class="hljs-keyword">None</span>, <span class="hljs-keyword">None</span>
    <span class="hljs-comment"># 过滤注释行</span>
    <span class="hljs-keyword">if</span> remainder.replace(<span class="hljs-string">' '</span>, <span class="hljs-string">''</span>)[<span class="hljs-number">0</span>] == <span class="hljs-string">'#'</span>:
        remainder = remainder.split(<span class="hljs-string">'\n'</span>, <span class="hljs-number">1</span>)[<span class="hljs-number">1</span>]

    <span class="hljs-comment"># 1空行视为分块的分隔符，这里读取head表示的分块</span>
    layer_param_block, remainder = remainder.split(<span class="hljs-string">'\n\n'</span>, <span class="hljs-number">1</span>)

    <span class="hljs-comment"># 处理得到的分块，并以'\n'将该块划分成行为元素的列表，等待处理</span>
    layer_param_lines = layer_param_block.split(<span class="hljs-string">'\n'</span>)[<span class="hljs-number">1</span>:]

    layer_name = str(self.layer_counter).zfill(<span class="hljs-number">3</span>) + <span class="hljs-string">'_'</span> + layer_type <span class="hljs-comment"># 当前块命名</span>
    layer_dict = dict(type=layer_type)

    <span class="hljs-comment"># 如果当前层是支持的，则进行处理，如yolo就不支持</span>
    <span class="hljs-keyword">if</span> layer_type <span class="hljs-keyword">in</span> self.supported_layers:
        <span class="hljs-keyword">for</span> param_line <span class="hljs-keyword">in</span> layer_param_lines:
            <span class="hljs-keyword">if</span> param_line[<span class="hljs-number">0</span>] == <span class="hljs-string">'#'</span>:
                <span class="hljs-keyword">continue</span>
            <span class="hljs-comment"># 解析每一行</span>
            param_type, param_value = self._parse_params(param_line)
            layer_dict[param_type] = param_value

    self.layer_counter += <span class="hljs-number">1</span>

    <span class="hljs-keyword">return</span> layer_dict, layer_name, remainder

<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">_parse_params</span><span class="hljs-params">(self, param_line)</span>:</span>
    <span class="hljs-string">"""解析每一行参数，当遇到layers时，返回list，其余返回字符串，整数，浮点数类型.

    Keyword argument:
    param_line -- 块中的一行需要解析的参数行
    """</span>
    param_line = param_line.replace(<span class="hljs-string">' '</span>, <span class="hljs-string">''</span>) <span class="hljs-comment"># 紧凑一下</span>
    param_type, param_value_raw = param_line.split(<span class="hljs-string">'='</span>) <span class="hljs-comment"># 以‘=’划分</span>
    param_value = <span class="hljs-keyword">None</span>

    <span class="hljs-comment"># 如果当前参数是layers，则以列表形式返回</span>
    <span class="hljs-keyword">if</span> param_type == <span class="hljs-string">'layers'</span>:
        layer_indexes = list()
        <span class="hljs-keyword">for</span> index <span class="hljs-keyword">in</span> param_value_raw.split(<span class="hljs-string">','</span>):
            layer_indexes.append(int(index))
        param_value = layer_indexes
    <span class="hljs-comment"># 否则先检测是否是整数，还是浮点数，不然就返回字符串类型</span>
    <span class="hljs-keyword">elif</span> isinstance(param_value_raw, str) <span class="hljs-keyword">and</span> <span class="hljs-keyword">not</span> param_value_raw.isalpha():
        condition_param_value_positive = param_value_raw.isdigit()
        condition_param_value_negative = param_value_raw[<span class="hljs-number">0</span>] == <span class="hljs-string">'-'</span> <span class="hljs-keyword">and</span> \
            param_value_raw[<span class="hljs-number">1</span>:].isdigit()
        <span class="hljs-keyword">if</span> condition_param_value_positive <span class="hljs-keyword">or</span> condition_param_value_negative:
            param_value = int(param_value_raw)
        <span class="hljs-keyword">else</span>:
            param_value = float(param_value_raw)
    <span class="hljs-keyword">else</span>:
        param_value = str(param_value_raw)

    <span class="hljs-keyword">return</span> param_type, param_value

‘’‘main第四步：被第三步类的_make_onnx_node方法调用 ‘’’
class MajorNodeSpecs(object):
“”“Helper class用于存储ONNX输出节点的信息，对应DarkNet 层的输出和该层输出通道，
一些DarkNet层并未被创建，因此没有对应的ONNX 节点，
不过仍然需要对其进行追踪以建立skip 连接
“””

<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">__init__</span><span class="hljs-params">(self, name, channels)</span>:</span>
    <span class="hljs-string">""" 初始化一个MajorNodeSpecs对象

    Keyword arguments:
    name -- ONNX节点的名称
    channels -- 该节点的输出通道的数量
    """</span>
    self.name = name
    self.channels = channels
    <span class="hljs-comment"># 对于yolov3.cfg中三层yolo层，这里表示该节点并非onnx节点，默认复制false</span>
    <span class="hljs-comment"># 其他如卷积，上采样等都是被赋予true</span>
    self.created_onnx_node = <span class="hljs-keyword">False</span>
    <span class="hljs-keyword">if</span> name <span class="hljs-keyword">is</span> <span class="hljs-keyword">not</span> <span class="hljs-keyword">None</span> <span class="hljs-keyword">and</span> isinstance(channels, int) <span class="hljs-keyword">and</span> channels &gt; <span class="hljs-number">0</span>:
        self.created_onnx_node = <span class="hljs-keyword">True</span>

‘’‘main第四步：被第三步类的_make_conv_node方法调用 ‘’’
class ConvParams(object):
“”"Helper class用于存储卷积层的超参数,包括在ONNX graph中的前置name和
为了卷积，偏置，BN等权重期望的维度

另外该类还扮演着为所有权重生成安全名称的封装，并检查合适的组合搭配
"""</span>

<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">__init__</span><span class="hljs-params">(self, node_name, batch_normalize, conv_weight_dims)</span>:</span>
    <span class="hljs-string">"""基于base 节点名称 (e.g. 101_convolutional),BN设置，卷积权重shape的构造器

    Keyword arguments:
    node_name -- YOLO卷积层的base名称
    batch_normalize -- bool值，表示是否使用BN
    conv_weight_dims -- 该层的卷积权重的维度
    """</span>
    self.node_name = node_name
    self.batch_normalize = batch_normalize
    <span class="hljs-keyword">assert</span> len(conv_weight_dims) == <span class="hljs-number">4</span>
    self.conv_weight_dims = conv_weight_dims

<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">generate_param_name</span><span class="hljs-params">(self, param_category, suffix)</span>:</span>
    <span class="hljs-string">"""基于两个字符串输入生成一个名称,并检查组合搭配是否合理"""</span>
    <span class="hljs-keyword">assert</span> suffix
    <span class="hljs-keyword">assert</span> param_category <span class="hljs-keyword">in</span> [<span class="hljs-string">'bn'</span>, <span class="hljs-string">'conv'</span>]
    <span class="hljs-keyword">assert</span>(suffix <span class="hljs-keyword">in</span> [<span class="hljs-string">'scale'</span>, <span class="hljs-string">'mean'</span>, <span class="hljs-string">'var'</span>, <span class="hljs-string">'weights'</span>, <span class="hljs-string">'bias'</span>])
    <span class="hljs-keyword">if</span> param_category == <span class="hljs-string">'bn'</span>:
        <span class="hljs-keyword">assert</span> self.batch_normalize
        <span class="hljs-keyword">assert</span> suffix <span class="hljs-keyword">in</span> [<span class="hljs-string">'scale'</span>, <span class="hljs-string">'bias'</span>, <span class="hljs-string">'mean'</span>, <span class="hljs-string">'var'</span>]
    <span class="hljs-keyword">elif</span> param_category == <span class="hljs-string">'conv'</span>:
        <span class="hljs-keyword">assert</span> suffix <span class="hljs-keyword">in</span> [<span class="hljs-string">'weights'</span>, <span class="hljs-string">'bias'</span>]
        <span class="hljs-keyword">if</span> suffix == <span class="hljs-string">'bias'</span>:
            <span class="hljs-keyword">assert</span> <span class="hljs-keyword">not</span> self.batch_normalize
    param_name = self.node_name + <span class="hljs-string">'_'</span> + param_category + <span class="hljs-string">'_'</span> + suffix
    <span class="hljs-keyword">return</span> param_name

‘’‘man第四步：被第三步类的build_onnx_graph方法调用 ‘’’
class WeightLoader(object):
“”“Helper class用于载入序列化的权重，
“””

<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">__init__</span><span class="hljs-params">(self, weights_file_path)</span>:</span>
    <span class="hljs-string">"""读取YOLOv3权重文件

    Keyword argument:
    weights_file_path --权重文件的路径.
    """</span>
    self.weights_file = self._open_weights_file(weights_file_path)

<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">load_conv_weights</span><span class="hljs-params">(self, conv_params)</span>:</span>
    <span class="hljs-string">"""返回权重文件的初始化器和卷积层的输入tensor

    Keyword argument:
    conv_params -- a ConvParams object
    """</span>
    initializer = list()
    inputs = list()
    <span class="hljs-keyword">if</span> conv_params.batch_normalize:
        <span class="hljs-comment"># 创建BN需要的bias，scale，mean，var等参数</span>
        bias_init, bias_input = self._create_param_tensors(
            conv_params, <span class="hljs-string">'bn'</span>, <span class="hljs-string">'bias'</span>)
        bn_scale_init, bn_scale_input = self._create_param_tensors(
            conv_params, <span class="hljs-string">'bn'</span>, <span class="hljs-string">'scale'</span>)
        bn_mean_init, bn_mean_input = self._create_param_tensors(
            conv_params, <span class="hljs-string">'bn'</span>, <span class="hljs-string">'mean'</span>)
        bn_var_init, bn_var_input = self._create_param_tensors(
            conv_params, <span class="hljs-string">'bn'</span>, <span class="hljs-string">'var'</span>)
        <span class="hljs-comment"># 初始化器扩展； 当前层输入的扩展</span>
        initializer.extend(
            [bn_scale_init, bias_init, bn_mean_init, bn_var_init])
        inputs.extend([bn_scale_input, bias_input,
                       bn_mean_input, bn_var_input])
    <span class="hljs-keyword">else</span>:
        <span class="hljs-comment"># 处理卷积层；  初始化器扩展； 当前层输入的扩展</span>
        bias_init, bias_input = self._create_param_tensors(
            conv_params, <span class="hljs-string">'conv'</span>, <span class="hljs-string">'bias'</span>)
        initializer.append(bias_init)
        inputs.append(bias_input)

    <span class="hljs-comment"># 创建卷积层权重；  初始化器扩展； 当前层输入的扩展</span>
    conv_init, conv_input = self._create_param_tensors(
        conv_params, <span class="hljs-string">'conv'</span>, <span class="hljs-string">'weights'</span>)
    initializer.append(conv_init)
    inputs.append(conv_input)

    <span class="hljs-keyword">return</span> initializer, inputs

<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">_open_weights_file</span><span class="hljs-params">(self, weights_file_path)</span>:</span>
    <span class="hljs-string">"""打开Yolov3 DarkNet文件流，并跳过开头.

    Keyword argument:
    weights_file_path -- 权重文件路径
    """</span>
    weights_file = open(weights_file_path, <span class="hljs-string">'rb'</span>)
    length_header = <span class="hljs-number">5</span>
    np.ndarray(
        shape=(length_header, ), dtype=<span class="hljs-string">'int32'</span>, buffer=weights_file.read(
            length_header * <span class="hljs-number">4</span>))
    <span class="hljs-keyword">return</span> weights_file

<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">_create_param_tensors</span><span class="hljs-params">(self, conv_params, param_category, suffix)</span>:</span>
    <span class="hljs-string">"""用权重文件中，与输入tensors一起的权重去初始化一个初始化器.

    Keyword arguments:
    conv_params -- a ConvParams object
    param_category -- the category of parameters to be created ('bn' or 'conv')
    suffix -- a string determining the sub-type of above param_category (e.g.,
    'weights' or 'bias')
    """</span>
    param_name, param_data, param_data_shape = self._load_one_param_type(
        conv_params, param_category, suffix)

    <span class="hljs-comment"># 调用onnx.helper.make_tensor</span>
    initializer_tensor = helper.make_tensor(
        param_name, TensorProto.FLOAT, param_data_shape, param_data)
    <span class="hljs-comment"># 调用onnx.helper.make_tensor_value_info</span>
    input_tensor = helper.make_tensor_value_info(
        param_name, TensorProto.FLOAT, param_data_shape)

    <span class="hljs-keyword">return</span> initializer_tensor, input_tensor

<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">_load_one_param_type</span><span class="hljs-params">(self, conv_params, param_category, suffix)</span>:</span>
    <span class="hljs-string">"""基于DarkNet顺序进行文件流的反序列化.

    Keyword arguments:
    conv_params -- a ConvParams object
    param_category -- the category of parameters to be created ('bn' or 'conv')
    suffix -- a string determining the sub-type of above param_category (e.g.,
    'weights' or 'bias')
    """</span>
    <span class="hljs-comment"># 生成合理的名称</span>
    param_name = conv_params.generate_param_name(param_category, suffix)
    channels_out, channels_in, filter_h, filter_w = conv_params.conv_weight_dims

    <span class="hljs-keyword">if</span> param_category == <span class="hljs-string">'bn'</span>:
        param_shape = [channels_out]
    <span class="hljs-keyword">elif</span> param_category == <span class="hljs-string">'conv'</span>:
        <span class="hljs-keyword">if</span> suffix == <span class="hljs-string">'weights'</span>:
            param_shape = [channels_out, channels_in, filter_h, filter_w]
        <span class="hljs-keyword">elif</span> suffix == <span class="hljs-string">'bias'</span>:
            param_shape = [channels_out]

    param_size = np.product(np.array(param_shape)) <span class="hljs-comment"># 计算参数的size</span>
    <span class="hljs-comment"># 用weights_file.read去逐字节的读取数据并转换</span>
    param_data = np.ndarray(
        shape=param_shape,
        dtype=<span class="hljs-string">'float32'</span>,
        buffer=self.weights_file.read(param_size * <span class="hljs-number">4</span>))
    param_data = param_data.flatten().astype(float)

    <span class="hljs-keyword">return</span> param_name, param_data, param_shape

‘’‘main第三步 ‘’’
class GraphBuilderONNX(object):
“”“用于创建ONNX graph的类，基于之前从yolov3.cfg读取的网络结构。该类函数方法有：
build_onnx_graph : 构建
_make_onnx_node
_make_input_tensor
_get_previous_node_specs
_make_conv_node
_make_shortcut_node
_make_route_node
_make_upsample_node
“””

<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">__init__</span><span class="hljs-params">(self, output_tensors)</span>:</span>
    <span class="hljs-string">"""用所有DarkNet默认参数来初始化；
        然后基于output_tensors指定输出维度；
       以他们的name为key

    Keyword argument:
    output_tensors -- 一个 OrderedDict类型
    """</span>

    self.output_tensors = output_tensors
    self._nodes = list()
    self.graph_def = <span class="hljs-keyword">None</span>
    self.input_tensor = <span class="hljs-keyword">None</span>
    self.epsilon_bn = <span class="hljs-number">1e-5</span>
    self.momentum_bn = <span class="hljs-number">0.99</span>
    self.alpha_lrelu = <span class="hljs-number">0.1</span>
    self.param_dict = OrderedDict()
    self.major_node_specs = list()
    self.batch_size = <span class="hljs-number">1</span>

<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">build_onnx_graph</span><span class="hljs-params">(
        self,
        layer_configs,
        weights_file_path,
        verbose=True)</span>:</span>
    <span class="hljs-string">"""基于所有的层配置进行迭代，创建一个ONNX graph，
        然后用下载的yolov3 权重文件进行填充，最后返回该graph定义.

    Keyword arguments:
    layer_configs -- OrderedDict对象，包含所有解析的层的配置
    weights_file_path -- 权重文件的位置
    verbose -- 是否在创建之后显示该graph(default: True)
    """</span>

    <span class="hljs-keyword">for</span> layer_name <span class="hljs-keyword">in</span> layer_configs.keys():

        layer_dict = layer_configs[layer_name]
        <span class="hljs-comment"># 读取yolov3.cfg中每一层，并将其作为onnx的节点</span>
        major_node_specs = self._make_onnx_node(layer_name, layer_dict)
        <span class="hljs-comment"># 如果当前为主要节点，则追加起来</span>
        <span class="hljs-keyword">if</span> major_node_specs.name:
            self.major_node_specs.append(major_node_specs)

    outputs = list()
    <span class="hljs-keyword">for</span> tensor_name <span class="hljs-keyword">in</span> self.output_tensors.keys():
        <span class="hljs-comment"># 将输出节点进行维度扩充</span>
        output_dims = [self.batch_size, ] + \
            self.output_tensors[tensor_name]
        <span class="hljs-comment"># 调用onnx的helper.make_tensor_value_info构建onnx张量，此时并未填充权重</span>
        output_tensor = helper.make_tensor_value_info(
            tensor_name, TensorProto.FLOAT, output_dims)
        outputs.append(output_tensor)

    inputs = [self.input_tensor]
    weight_loader = WeightLoader(weights_file_path)
    initializer = list()
    <span class="hljs-comment"># self.param_dict在_make_onnx_node中已处理</span>
    <span class="hljs-keyword">for</span> layer_name <span class="hljs-keyword">in</span> self.param_dict.keys():
        _, layer_type = layer_name.split(<span class="hljs-string">'_'</span>, <span class="hljs-number">1</span>) <span class="hljs-comment"># 如001_convolutional</span>
        conv_params = self.param_dict[layer_name]
        <span class="hljs-keyword">assert</span> layer_type == <span class="hljs-string">'convolutional'</span>
        initializer_layer, inputs_layer = weight_loader.load_conv_weights(
            conv_params)
        initializer.extend(initializer_layer)
        inputs.extend(inputs_layer)
    <span class="hljs-keyword">del</span> weight_loader

    <span class="hljs-comment"># 调用onnx的helper.make_graph进行onnx graph的构建</span>
    self.graph_def = helper.make_graph(
        nodes=self._nodes,
        name=<span class="hljs-string">'YOLOv3-608'</span>,
        inputs=inputs,
        outputs=outputs,
        initializer=initializer
    )

    <span class="hljs-keyword">if</span> verbose:
        print(helper.printable_graph(self.graph_def))

    <span class="hljs-comment"># 调用onnx的helper.make_model进行模型的构建</span>
    model_def = helper.make_model(self.graph_def,
                                  producer_name=<span class="hljs-string">'NVIDIA TensorRT sample'</span>)
    <span class="hljs-keyword">return</span> model_def

<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">_make_onnx_node</span><span class="hljs-params">(self, layer_name, layer_dict)</span>:</span>
    <span class="hljs-string">"""输入一个layer参数字典，选择对应的函数来创建ONNX节点，然后将为图创建的重要的信息存储为
       MajorNodeSpec对象

    Keyword arguments:
    layer_name -- layer的名称 (即layer_configs中的key)
    layer_dict -- 一个layer参数字典 (layer_configs的value)
    """</span>
    
    layer_type = layer_dict[<span class="hljs-string">'type'</span>]
    <span class="hljs-comment"># 先检查self.input_tensor是否为空，为空且第一个块不是net，则报错，否则处理该net</span>
    <span class="hljs-comment"># 可以看出 这里只在最开始执行一次，因为后续self.input_tensor都不为空。</span>
    <span class="hljs-keyword">if</span> self.input_tensor <span class="hljs-keyword">is</span> <span class="hljs-keyword">None</span>:
        <span class="hljs-keyword">if</span> layer_type == <span class="hljs-string">'net'</span>:
            major_node_output_name, major_node_output_channels = self._make_input_tensor(
                layer_name, layer_dict)
            major_node_specs = MajorNodeSpecs(major_node_output_name,
                                              major_node_output_channels)
        <span class="hljs-keyword">else</span>:
            <span class="hljs-keyword">raise</span> ValueError(<span class="hljs-string">'The first node has to be of type "net".'</span>)
    <span class="hljs-keyword">else</span>:
        node_creators = dict()
        node_creators[<span class="hljs-string">'convolutional'</span>] = self._make_conv_node
        node_creators[<span class="hljs-string">'shortcut'</span>] = self._make_shortcut_node
        node_creators[<span class="hljs-string">'route'</span>] = self._make_route_node
        node_creators[<span class="hljs-string">'upsample'</span>] = self._make_upsample_node

        <span class="hljs-comment"># 依次处理不同的层，并调用对应node_creators[layer_type]()函数进行处理</span>
        <span class="hljs-keyword">if</span> layer_type <span class="hljs-keyword">in</span> node_creators.keys():
            major_node_output_name, major_node_output_channels = \
                node_creators[layer_type](layer_name, layer_dict)
            major_node_specs = MajorNodeSpecs(major_node_output_name,
                                              major_node_output_channels)
        <span class="hljs-keyword">else</span>:
            <span class="hljs-comment"># 跳过三个yolo层</span>
            print(
                <span class="hljs-string">'Layer of type %s not supported, skipping ONNX node generation.'</span> %
                layer_type)
            major_node_specs = MajorNodeSpecs(layer_name,
                                              <span class="hljs-keyword">None</span>)
    <span class="hljs-keyword">return</span> major_node_specs

<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">_make_input_tensor</span><span class="hljs-params">(self, layer_name, layer_dict)</span>:</span>
    <span class="hljs-string">"""为net layer创建输入tensor，并存储对应batch size.可以看出，该函数只被调用一次

    Keyword arguments:
    layer_name -- 层的名字 (如 layer_configs中key)
    layer_dict -- 一个layer参数字典( layer_configs中的value)
    """</span>
    batch_size = layer_dict[<span class="hljs-string">'batch'</span>]
    channels = layer_dict[<span class="hljs-string">'channels'</span>]
    height = layer_dict[<span class="hljs-string">'height'</span>]
    width = layer_dict[<span class="hljs-string">'width'</span>]
    self.batch_size = batch_size

    <span class="hljs-comment"># 用onnx.helper.make_tensor_value_info构建onnx张量节点</span>
    input_tensor = helper.make_tensor_value_info(
        str(layer_name), TensorProto.FLOAT, [
            batch_size, channels, height, width])
    self.input_tensor = input_tensor

    <span class="hljs-keyword">return</span> layer_name, channels

<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">_get_previous_node_specs</span><span class="hljs-params">(self, target_index=<span class="hljs-number">-1</span>)</span>:</span>
    <span class="hljs-string">"""获取之前创建好的onnx节点(跳过那些没生成的节点，比如yolo节点).
    target_index可以能够直接跳到对应节点.

    Keyword arguments:
    target_index -- 可选的参数，帮助跳到具体索引(default: -1 表示跳到前一个元素)
    """</span>

    <span class="hljs-comment"># 通过反向遍历，找到最后一个（这里是第一个）created_onnx_node为真的节点</span>
    previous_node = <span class="hljs-keyword">None</span>
    <span class="hljs-keyword">for</span> node <span class="hljs-keyword">in</span> self.major_node_specs[target_index::<span class="hljs-number">-1</span>]:
        <span class="hljs-keyword">if</span> node.created_onnx_node:
            previous_node = node
            <span class="hljs-keyword">break</span>
    <span class="hljs-keyword">assert</span> previous_node <span class="hljs-keyword">is</span> <span class="hljs-keyword">not</span> <span class="hljs-keyword">None</span>
    <span class="hljs-keyword">return</span> previous_node

<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">_make_conv_node</span><span class="hljs-params">(self, layer_name, layer_dict)</span>:</span>
    <span class="hljs-string">"""用可选的bn和激活函数nonde去创建一个onnx的卷积node

    Keyword arguments:
    layer_name -- 层的名字 (如 layer_configs中key)
    layer_dict -- 一个layer参数字典( layer_configs中的value)
    """</span>
    <span class="hljs-comment"># 先找最近的一个节点</span>
    previous_node_specs = self._get_previous_node_specs()

    <span class="hljs-string">''' i) 处理卷积层'''</span>
    <span class="hljs-comment"># 构建该层的inputs，通道等等信息</span>
    inputs = [previous_node_specs.name]
    previous_channels = previous_node_specs.channels
    kernel_size = layer_dict[<span class="hljs-string">'size'</span>]
    stride = layer_dict[<span class="hljs-string">'stride'</span>]
    filters = layer_dict[<span class="hljs-string">'filters'</span>]
    <span class="hljs-comment"># 检测该层是否有bn</span>
    batch_normalize = <span class="hljs-keyword">False</span>
    <span class="hljs-keyword">if</span> <span class="hljs-string">'batch_normalize'</span> <span class="hljs-keyword">in</span> layer_dict.keys(
    ) <span class="hljs-keyword">and</span> layer_dict[<span class="hljs-string">'batch_normalize'</span>] == <span class="hljs-number">1</span>:
        batch_normalize = <span class="hljs-keyword">True</span>

    kernel_shape = [kernel_size, kernel_size]
    weights_shape = [filters, previous_channels] + kernel_shape
    <span class="hljs-comment"># 构建卷积层的参数层的实例</span>
    conv_params = ConvParams(layer_name, batch_normalize, weights_shape)

    strides = [stride, stride]
    dilations = [<span class="hljs-number">1</span>, <span class="hljs-number">1</span>]
    <span class="hljs-comment"># 调用ConvParams.generate_param_name生成合适的参数名称</span>
    weights_name = conv_params.generate_param_name(<span class="hljs-string">'conv'</span>, <span class="hljs-string">'weights'</span>)
    inputs.append(weights_name)
    <span class="hljs-keyword">if</span> <span class="hljs-keyword">not</span> batch_normalize:
        bias_name = conv_params.generate_param_name(<span class="hljs-string">'conv'</span>, <span class="hljs-string">'bias'</span>)
        inputs.append(bias_name)

    <span class="hljs-comment"># 用onnx.helper.make_node构建onnx的卷积节点</span>
    conv_node = helper.make_node(
        <span class="hljs-string">'Conv'</span>,
        inputs=inputs,
        outputs=[layer_name],
        kernel_shape=kernel_shape,
        strides=strides,
        auto_pad=<span class="hljs-string">'SAME_LOWER'</span>,
        dilations=dilations,
        name=layer_name
    )
    self._nodes.append(conv_node)

    inputs = [layer_name]
    layer_name_output = layer_name

    <span class="hljs-string">''' ii) 处理BN层'''</span>
    <span class="hljs-keyword">if</span> batch_normalize:
        layer_name_bn = layer_name + <span class="hljs-string">'_bn'</span>
        bn_param_suffixes = [<span class="hljs-string">'scale'</span>, <span class="hljs-string">'bias'</span>, <span class="hljs-string">'mean'</span>, <span class="hljs-string">'var'</span>]
        <span class="hljs-keyword">for</span> suffix <span class="hljs-keyword">in</span> bn_param_suffixes:
            bn_param_name = conv_params.generate_param_name(<span class="hljs-string">'bn'</span>, suffix)
            inputs.append(bn_param_name)
        batchnorm_node = helper.make_node(
            <span class="hljs-string">'BatchNormalization'</span>,
            inputs=inputs,
            outputs=[layer_name_bn],
            epsilon=self.epsilon_bn,
            momentum=self.momentum_bn,
            name=layer_name_bn
        )
        self._nodes.append(batchnorm_node)

        inputs = [layer_name_bn]
        layer_name_output = layer_name_bn

    <span class="hljs-string">''' iii) 处理激活函数'''</span>
    <span class="hljs-keyword">if</span> layer_dict[<span class="hljs-string">'activation'</span>] == <span class="hljs-string">'leaky'</span>:
        layer_name_lrelu = layer_name + <span class="hljs-string">'_lrelu'</span>

        lrelu_node = helper.make_node(
            <span class="hljs-string">'LeakyRelu'</span>,
            inputs=inputs,
            outputs=[layer_name_lrelu],
            name=layer_name_lrelu,
            alpha=self.alpha_lrelu
        )
        self._nodes.append(lrelu_node)
        inputs = [layer_name_lrelu]
        layer_name_output = layer_name_lrelu
    <span class="hljs-keyword">elif</span> layer_dict[<span class="hljs-string">'activation'</span>] == <span class="hljs-string">'linear'</span>:
        <span class="hljs-keyword">pass</span>
    <span class="hljs-keyword">else</span>:
        print(<span class="hljs-string">'Activation not supported.'</span>)

    self.param_dict[layer_name] = conv_params
    <span class="hljs-keyword">return</span> layer_name_output, filters

<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">_make_shortcut_node</span><span class="hljs-params">(self, layer_name, layer_dict)</span>:</span>
    <span class="hljs-string">"""从DarkNet graph中读取信息，基于onnx 的add 节点创建shortcut 节点.

    Keyword arguments:
    layer_name -- 层的名字 (如 layer_configs中key)
    layer_dict -- 一个layer参数字典( layer_configs中的value)
    """</span>
    shortcut_index = layer_dict[<span class="hljs-string">'from'</span>] <span class="hljs-comment"># 当前层与前面哪层shorcut</span>
    activation = layer_dict[<span class="hljs-string">'activation'</span>]
    <span class="hljs-keyword">assert</span> activation == <span class="hljs-string">'linear'</span>

    first_node_specs = self._get_previous_node_specs() <span class="hljs-comment"># 最近一层</span>
    second_node_specs = self._get_previous_node_specs(
        target_index=shortcut_index) <span class="hljs-comment"># 前面具体需要shorcut的层</span>
    <span class="hljs-keyword">assert</span> first_node_specs.channels == second_node_specs.channels
    channels = first_node_specs.channels
    inputs = [first_node_specs.name, second_node_specs.name]
    <span class="hljs-comment"># 用onnx.helper.make_node创建节点</span>
    shortcut_node = helper.make_node(
        <span class="hljs-string">'Add'</span>,
        inputs=inputs,
        outputs=[layer_name],
        name=layer_name,
    )
    self._nodes.append(shortcut_node)
    <span class="hljs-keyword">return</span> layer_name, channels

<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">_make_route_node</span><span class="hljs-params">(self, layer_name, layer_dict)</span>:</span>
    <span class="hljs-string">"""如果来自DarkNet配置的layer参数只有一个所以，那么接着在指定（负）索引上创建节点
       否则，创建一个onnx concat 节点来实现路由特性.

    Keyword arguments:
    layer_name -- 层的名字 (如 layer_configs中key)
    layer_dict -- 一个layer参数字典( layer_configs中的value)
    """</span>
    <span class="hljs-comment"># 处理yolov3.cfg中[route]</span>
    route_node_indexes = layer_dict[<span class="hljs-string">'layers'</span>]

    <span class="hljs-keyword">if</span> len(route_node_indexes) == <span class="hljs-number">1</span>:
        split_index = route_node_indexes[<span class="hljs-number">0</span>]
        <span class="hljs-keyword">assert</span> split_index &lt; <span class="hljs-number">0</span>
        <span class="hljs-comment"># Increment by one because we skipped the YOLO layer:</span>
        split_index += <span class="hljs-number">1</span>
        self.major_node_specs = self.major_node_specs[:split_index]
        layer_name = <span class="hljs-keyword">None</span>
        channels = <span class="hljs-keyword">None</span>
    <span class="hljs-keyword">else</span>:
        inputs = list()
        channels = <span class="hljs-number">0</span>
        <span class="hljs-keyword">for</span> index <span class="hljs-keyword">in</span> route_node_indexes:
            <span class="hljs-keyword">if</span> index &gt; <span class="hljs-number">0</span>:
                <span class="hljs-comment"># Increment by one because we count the input as a node (DarkNet</span>
                <span class="hljs-comment"># does not)</span>
                index += <span class="hljs-number">1</span>
            route_node_specs = self._get_previous_node_specs(
                target_index=index)
            inputs.append(route_node_specs.name)
            channels += route_node_specs.channels
        <span class="hljs-keyword">assert</span> inputs
        <span class="hljs-keyword">assert</span> channels &gt; <span class="hljs-number">0</span>

        route_node = helper.make_node(
            <span class="hljs-string">'Concat'</span>,
            axis=<span class="hljs-number">1</span>,
            inputs=inputs,
            outputs=[layer_name],
            name=layer_name,
        )
        self._nodes.append(route_node)
    <span class="hljs-keyword">return</span> layer_name, channels

<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">_make_upsample_node</span><span class="hljs-params">(self, layer_name, layer_dict)</span>:</span>
    <span class="hljs-string">"""创建一个onnx的Upsample节点.

    Keyword arguments:
    layer_name -- 层的名字 (如 layer_configs中key)
    layer_dict -- 一个layer参数字典( layer_configs中的value)
    """</span>
    upsample_factor = float(layer_dict[<span class="hljs-string">'stride'</span>])
    previous_node_specs = self._get_previous_node_specs()
    inputs = [previous_node_specs.name]
    channels = previous_node_specs.channels
    <span class="hljs-keyword">assert</span> channels &gt; <span class="hljs-number">0</span>
    upsample_node = helper.make_node(
        <span class="hljs-string">'Upsample'</span>,
        mode=<span class="hljs-string">'nearest'</span>,
        <span class="hljs-comment"># For ONNX versions &lt;0.7.0, Upsample nodes accept different parameters than 'scales':</span>
        scales=[<span class="hljs-number">1.0</span>, <span class="hljs-number">1.0</span>, upsample_factor, upsample_factor],
        inputs=inputs,
        outputs=[layer_name],
        name=layer_name,
    )
    self._nodes.append(upsample_node)
    <span class="hljs-keyword">return</span> layer_name, channels

def generate_md5_checksum(local_path):
“”"计算本地文件的md5

Keyword argument:
local_path -- 本地文件路径
"""</span>
<span class="hljs-keyword">with</span> open(local_path) <span class="hljs-keyword">as</span> local_file:
    data = local_file.read()
    <span class="hljs-keyword">return</span> hashlib.md5(data).hexdigest()

def download_file(local_path, link, checksum_reference=None):
“”"下载指定url到本地，并进行摘要校对.

Keyword arguments:
local_path -- 本地文件存储路径
link -- 需要下载的url
checksum_reference -- expected MD5 checksum of the file
"""</span>
<span class="hljs-keyword">if</span> <span class="hljs-keyword">not</span> os.path.exists(local_path):
    print(<span class="hljs-string">'Downloading from %s, this may take a while...'</span> % link)
    wget.download(link, local_path)
    print()

<span class="hljs-keyword">if</span> checksum_reference <span class="hljs-keyword">is</span> <span class="hljs-keyword">not</span> <span class="hljs-keyword">None</span>:
    checksum = generate_md5_checksum(local_path)
    <span class="hljs-keyword">if</span> checksum != checksum_reference:
        <span class="hljs-keyword">raise</span> ValueError(
            <span class="hljs-string">'The MD5 checksum of local file %s differs from %s, please manually remove \
             the file and try again.'</span> %
            (local_path, checksum_reference))

<span class="hljs-keyword">return</span> local_path

def main():

<span class="hljs-string">"""Run the DarkNet-to-ONNX conversion for YOLOv3-608."""</span>

<span class="hljs-comment"># 注释掉下面的部分，</span>

# if sys.version_info[0] > 2:
# raise Exception(“This is script is only compatible with python2, please re-run this script </span>
# with python2. The rest of this sample can be run with either version of python”)

<span class="hljs-string">''' 1 - 下载yolov3的配置文件，并进行摘要验证'''</span>
cfg_file_path = download_file(
    <span class="hljs-string">'yolov3.cfg'</span>,

‘https://raw.githubusercontent.com/pjreddie/darknet/f86901f6177dfc6116360a13cc06ab680e0c86b0/cfg/yolov3.cfg’,
‘b969a43a848bbf26901643b833cfb96c’)

<span class="hljs-comment"># DarkNetParser将会只提取这些层的参数，类型为'yolo'的这三层不能很好的解析，</span>
<span class="hljs-comment"># 因为他们包含在后续的后处理中；</span>
supported_layers = [<span class="hljs-string">'net'</span>, <span class="hljs-string">'convolutional'</span>, <span class="hljs-string">'shortcut'</span>,
                    <span class="hljs-string">'route'</span>, <span class="hljs-string">'upsample'</span>]

<span class="hljs-string">''' 2 - 创建一个DarkNetParser对象，并生成一个OrderedDict，包含cfg文件读取的所有层配置'''</span>
parser = DarkNetParser(supported_layers)
layer_configs = parser.parse_cfg_file(cfg_file_path)
<span class="hljs-comment"># 在解析完之后，不再需要该对象</span>
<span class="hljs-keyword">del</span> parser

<span class="hljs-string">''' 3 - 实例化一个GraphBuilderONNX类对象，用已知输出tensor维度进行初始化'''</span>
<span class="hljs-comment"># 在上面的layer_config，有三个输出是需要知道的，CHW格式</span>
output_tensor_dims = OrderedDict()
output_tensor_dims[<span class="hljs-string">'082_convolutional'</span>] = [<span class="hljs-number">255</span>, <span class="hljs-number">19</span>, <span class="hljs-number">19</span>]
output_tensor_dims[<span class="hljs-string">'094_convolutional'</span>] = [<span class="hljs-number">255</span>, <span class="hljs-number">38</span>, <span class="hljs-number">38</span>]
output_tensor_dims[<span class="hljs-string">'106_convolutional'</span>] = [<span class="hljs-number">255</span>, <span class="hljs-number">76</span>, <span class="hljs-number">76</span>]

<span class="hljs-comment"># 内置yolov3的一些默认参数来进行实例化</span>
builder = GraphBuilderONNX(output_tensor_dims)

<span class="hljs-string">''' 4 - 调用GraphBuilderONNX的build_onnx_graph方法
       用之前解析好的层配置信息和权重文件，生成ONNX graph'''</span>
<span class="hljs-string">''' 从作者官网下载yolov3的权重文件，以此填充tensorrt的network '''</span>
weights_file_path = download_file(
    <span class="hljs-string">'yolov3.weights'</span>,
    <span class="hljs-string">'https://pjreddie.com/media/files/yolov3.weights'</span>,
    <span class="hljs-string">'c84e5b99d0e52cd466ae710cadf6d84c'</span>)

yolov3_model_def = builder.build_onnx_graph(
    layer_configs=layer_configs,
    weights_file_path=weights_file_path,
    verbose=<span class="hljs-keyword">True</span>)
<span class="hljs-comment"># 模型定义结束，删除builder对象</span>
<span class="hljs-keyword">del</span> builder

<span class="hljs-string">''' 5 - 在ONNX模型定义上进行健全检查'''</span>
onnx.checker.check_model(yolov3_model_def)

<span class="hljs-string">''' 6 - 序列化生成的ONNX graph到文件'''</span>
output_file_path = <span class="hljs-string">'yolov3.onnx'</span>
onnx.save(yolov3_model_def, output_file_path)

if name == ‘main’:
main()

结果如下：

[root@30d4bceec4c4 yolov3_onnx]# python yolov3_to_onnx.py
Layer of type yolo not supported, skipping ONNX node generation.
Layer of type yolo not supported, skipping ONNX node generation.
Layer of type yolo not supported, skipping ONNX node generation.
graph YOLOv3-608 (
  %000_net[FLOAT, 64x3x608x608]
) initializers (
  %001_convolutional_bn_scale[FLOAT, 32]
  %001_convolutional_bn_bias[FLOAT, 32]
  %001_convolutional_bn_mean[FLOAT, 32]
  %001_convolutional_bn_var[FLOAT, 32]
  %001_convolutional_conv_weights[FLOAT, 32x3x3x3]
  %002_convolutional_bn_scale[FLOAT, 64]
 ......
  %105_convolutional_conv_weights[FLOAT, 256x128x3x3]
  %106_convolutional_conv_bias[FLOAT, 255]
  %106_convolutional_conv_weights[FLOAT, 255x256x1x1]
) {
  %001_convolutional = Conv[auto_pad = 'SAME_LOWER', dilations = [1, 1], kernel_shape = [3, 3], strides = [1, 1]](%000_net, %001_convolutional_conv_weights)
  %001_convolutional_bn = BatchNormalization[epsilon = 9.99999974737875e-06, momentum = 0.990000009536743](%001_convolutional, 
......
  %105_convolutional_bn = BatchNormalization[epsilon = 9.99999974737875e-06, momentum = 0.990000009536743](%105_convolutional, %105_convolutional_bn_scale, %105_convolutional_bn_bias, %105_convolutional_bn_mean, %105_convolutional_bn_var)
  %105_convolutional_lrelu = LeakyRelu[alpha = 0.100000001490116](%105_convolutional_bn)
  %106_convolutional = Conv[auto_pad = 'SAME_LOWER', dilations = [1, 1], kernel_shape = [1, 1], strides = [1, 1]](%105_convolutional_lrelu, %106_convolutional_conv_weights, %106_convolutional_conv_bias)
  return %082_convolutional, %094_convolutional, %106_convolutional
}

ps:在该例子中onnx不要安装1.4.1版本，可以安装如1.2.1版本，否则会出现

此时会生成文件yolov3.onnx。

3 onnx转trt并进行inference

接下来看onnx_to_tensorrt.py

from __future__ import print_function

import numpy as np
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
from PIL import ImageDraw

from yolov3_to_onnx import download_file
from data_processing import PreprocessYOLO, PostprocessYOLO, ALL_CATEGORIES

import sys, os
#sys.path.insert(1, os.path.join(sys.path[0], “…”))
#import common

‘’‘main第3.2步 ‘’’
def allocate_buffers(engine):

inputs = []
outputs = []
bindings = []
stream = cuda.Stream()

<span class="hljs-keyword">for</span> binding <span class="hljs-keyword">in</span> engine:

    size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
    dtype = trt.nptype(engine.get_binding_dtype(binding))

    <span class="hljs-comment"># 分配host和device端的buffer</span>
    host_mem = cuda.pagelocked_empty(size, dtype)
    device_mem = cuda.mem_alloc(host_mem.nbytes)

    <span class="hljs-comment"># 将device端的buffer追加到device的bindings.</span>
    bindings.append(int(device_mem))

    <span class="hljs-comment"># Append to the appropriate list.</span>
    <span class="hljs-keyword">if</span> engine.binding_is_input(binding):
        inputs.append(HostDeviceMem(host_mem, device_mem))
    <span class="hljs-keyword">else</span>:
        outputs.append(HostDeviceMem(host_mem, device_mem))

<span class="hljs-keyword">return</span> inputs, outputs, bindings, stream

‘’‘main中第3.3步 ‘’’
# 该函数可以适应多个输入/输出；输入和输出格式为HostDeviceMem对象组成的列表
def do_inference(context, bindings, inputs, outputs, stream, batch_size=1):

<span class="hljs-comment"># 将数据移动到GPU</span>
[cuda.memcpy_htod_async(inp.device, inp.host, stream) <span class="hljs-keyword">for</span> inp <span class="hljs-keyword">in</span> inputs]

<span class="hljs-comment"># 执行inference.</span>
context.execute_async(batch_size=batch_size, bindings=bindings, stream_handle=stream.handle)

<span class="hljs-comment"># 将结果从 GPU写回到host端</span>
[cuda.memcpy_dtoh_async(out.host, out.device, stream) <span class="hljs-keyword">for</span> out <span class="hljs-keyword">in</span> outputs]

<span class="hljs-comment"># 同步stream</span>
stream.synchronize()

<span class="hljs-comment"># 返回host端的输出结果</span>
<span class="hljs-keyword">return</span> [out.host <span class="hljs-keyword">for</span> out <span class="hljs-keyword">in</span> outputs]

#------------

TRT_LOGGER = trt.Logger()

def draw_bboxes(image_raw, bboxes, confidences, categories, all_categories, bbox_color=‘blue’):
“”"在原始输入图片上标记bounding box没然后返回结果.

Keyword arguments:
image_raw -- a raw PIL Image
bboxes -- NumPy array containing the bounding box coordinates of N objects, with shape (N,4).
categories -- NumPy array containing the corresponding category for each object,
with shape (N,)
confidences -- NumPy array containing the corresponding confidence for each object,
with shape (N,)
all_categories -- a list of all categories in the correct ordered (required for looking up
the category name)
bbox_color -- an optional string specifying the color of the bounding boxes (default: 'blue')
"""</span>

draw = ImageDraw.Draw(image_raw)
print(bboxes, confidences, categories)
<span class="hljs-keyword">for</span> box, score, category <span class="hljs-keyword">in</span> zip(bboxes, confidences, categories):
    x_coord, y_coord, width, height = box
    left = max(<span class="hljs-number">0</span>, np.floor(x_coord + <span class="hljs-number">0.5</span>).astype(int))
    top = max(<span class="hljs-number">0</span>, np.floor(y_coord + <span class="hljs-number">0.5</span>).astype(int))
    right = min(image_raw.width, np.floor(x_coord + width + <span class="hljs-number">0.5</span>).astype(int))
    bottom = min(image_raw.height, np.floor(y_coord + height + <span class="hljs-number">0.5</span>).astype(int))

    draw.rectangle(((left, top), (right, bottom)), outline=bbox_color)
    draw.text((left, top - <span class="hljs-number">12</span>), <span class="hljs-string">'{0} {1:.2f}'</span>.format(all_categories[category], score), fill=bbox_color)

<span class="hljs-keyword">return</span> image_raw

def get_engine(onnx_file_path, engine_file_path=""):
“”“如果已经有序列化engine，则直接用，否则构建新的tensorrt engine然后保存.”""
# 闭包
def build_engine():
“”“Takes an ONNX file and creates a TensorRT engine to run inference with”""
with trt.Builder(TRT_LOGGER) as builder,
builder.create_network() as network,
trt.OnnxParser(network, TRT_LOGGER) as parser:

        builder.max_workspace_size = <span class="hljs-number">1</span> &lt;&lt; <span class="hljs-number">30</span> <span class="hljs-comment"># 1GB</span>
        builder.max_batch_size = <span class="hljs-number">1</span>

        <span class="hljs-comment"># 解析模型文件</span>
        <span class="hljs-keyword">if</span> <span class="hljs-keyword">not</span> os.path.exists(onnx_file_path):
            print(<span class="hljs-string">'ONNX file {} not found, please run yolov3_to_onnx.py first to generate it.'</span>.format(onnx_file_path))
            exit(<span class="hljs-number">0</span>)

        print(<span class="hljs-string">'Loading ONNX file from path {}...'</span>.format(onnx_file_path))
        <span class="hljs-keyword">with</span> open(onnx_file_path, <span class="hljs-string">'rb'</span>) <span class="hljs-keyword">as</span> model:
            print(<span class="hljs-string">'Beginning ONNX file parsing'</span>)
            parser.parse(model.read())
        print(<span class="hljs-string">'Completed parsing of ONNX file'</span>)

        print(<span class="hljs-string">'Building an engine from file {}; this may take a while...'</span>.format(onnx_file_path))
        engine = builder.build_cuda_engine(network)
        print(<span class="hljs-string">"Completed creating Engine"</span>)

        <span class="hljs-keyword">with</span> open(engine_file_path, <span class="hljs-string">"wb"</span>) <span class="hljs-keyword">as</span> f:
            f.write(engine.serialize())
        <span class="hljs-keyword">return</span> engine

<span class="hljs-keyword">if</span> os.path.exists(engine_file_path):
    <span class="hljs-comment"># 如果序列化engine已经存在，那么就直接跳过构建部分.</span>
    print(<span class="hljs-string">"Reading engine from file {}"</span>.format(engine_file_path))
    <span class="hljs-keyword">with</span> open(engine_file_path, <span class="hljs-string">"rb"</span>) <span class="hljs-keyword">as</span> f, \
         trt.Runtime(TRT_LOGGER) <span class="hljs-keyword">as</span> runtime:
        <span class="hljs-keyword">return</span> runtime.deserialize_cuda_engine(f.read())
<span class="hljs-keyword">else</span>:
    <span class="hljs-keyword">return</span> build_engine()

def main():
“”“Create a TensorRT engine for ONNX-based YOLOv3-608 and run inference.”""

<span class="hljs-string">''' 1 - 装载之前转换好的onnx，准备测试图片'''</span>
onnx_file_path = <span class="hljs-string">'yolov3.onnx'</span>
engine_file_path = <span class="hljs-string">"yolov3.trt"</span>
<span class="hljs-comment"># 下载测试图片</span>
input_image_path = download_file(<span class="hljs-string">'dog.jpg'</span>,
    <span class="hljs-string">'https://github.com/pjreddie/darknet/raw/f86901f6177dfc6116360a13cc06ab680e0c86b0/data/dog.jpg'</span>, checksum_reference=<span class="hljs-keyword">None</span>)

<span class="hljs-string">''' 2 - 对图片进行预处理'''</span>
<span class="hljs-comment"># yolov3网络的输入size，HW顺序</span>
input_resolution_yolov3_HW = (<span class="hljs-number">608</span>, <span class="hljs-number">608</span>)
<span class="hljs-comment"># 创建一个预处理来处理任意图片，以符合yolov3的输入</span>
preprocessor = PreprocessYOLO(input_resolution_yolov3_HW)
<span class="hljs-comment"># 载入图像，并进行预处理</span>
image_raw, image = preprocessor.process(input_image_path)
<span class="hljs-comment"># 将该预处理好的图像以WH格式存储，以备后续使用</span>
shape_orig_WH = image_raw.size

<span class="hljs-string">''' 3 - 基于tensorrt进行yolov3模型的运行'''</span>
<span class="hljs-comment"># yolov3输出的三个map的shapeOutput shapes expected by the post-processor</span>
output_shapes = [(<span class="hljs-number">1</span>, <span class="hljs-number">255</span>, <span class="hljs-number">19</span>, <span class="hljs-number">19</span>), (<span class="hljs-number">1</span>, <span class="hljs-number">255</span>, <span class="hljs-number">38</span>, <span class="hljs-number">38</span>), (<span class="hljs-number">1</span>, <span class="hljs-number">255</span>, <span class="hljs-number">76</span>, <span class="hljs-number">76</span>)]
<span class="hljs-comment"># 用 TensorRT进行inference</span>
trt_outputs = []

‘’’ 3.1 - 基于get_engine生成engine’’’
with get_engine(onnx_file_path, engine_file_path) as engine,
engine.create_execution_context() as context:

    <span class="hljs-string">''' 3.2 - 分配host，device端的buffer'''</span>
    inputs, outputs, bindings, stream = allocate_buffers(engine)
    
    print(<span class="hljs-string">'Running inference on image {}...'</span>.format(input_image_path))

    <span class="hljs-string">''' 3.3 - 进行inference'''</span>
    inputs[<span class="hljs-number">0</span>].host = image
    trt_outputs = do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)

<span class="hljs-string">''' 4 - 对tensorrt在onnx结构的yolov3上得到的结果进行后处理'''</span>
trt_outputs = [output.reshape(shape) <span class="hljs-keyword">for</span> output, shape <span class="hljs-keyword">in</span> zip(trt_outputs, output_shapes)]

postprocessor_args = {<span class="hljs-string">"yolo_masks"</span>: [(<span class="hljs-number">6</span>, <span class="hljs-number">7</span>, <span class="hljs-number">8</span>), (<span class="hljs-number">3</span>, <span class="hljs-number">4</span>, <span class="hljs-number">5</span>), (<span class="hljs-number">0</span>, <span class="hljs-number">1</span>, <span class="hljs-number">2</span>)],                    
                      <span class="hljs-string">"yolo_anchors"</span>: [(<span class="hljs-number">10</span>, <span class="hljs-number">13</span>), (<span class="hljs-number">16</span>, <span class="hljs-number">30</span>), (<span class="hljs-number">33</span>, <span class="hljs-number">23</span>), (<span class="hljs-number">30</span>, <span class="hljs-number">61</span>), (<span class="hljs-number">62</span>, <span class="hljs-number">45</span>), 
                                       (<span class="hljs-number">59</span>, <span class="hljs-number">119</span>), (<span class="hljs-number">116</span>, <span class="hljs-number">90</span>), (<span class="hljs-number">156</span>, <span class="hljs-number">198</span>), (<span class="hljs-number">373</span>, <span class="hljs-number">326</span>)],
                      <span class="hljs-string">"obj_threshold"</span>: <span class="hljs-number">0.6</span>,         <span class="hljs-comment"># 对象覆盖的阈值，[0,1]之间</span>
                      <span class="hljs-string">"nms_threshold"</span>: <span class="hljs-number">0.5</span>,       <span class="hljs-comment"># nms的阈值，[0,1]之间</span>
                      <span class="hljs-string">"yolo_input_resolution"</span>: input_resolution_yolov3_HW}
<span class="hljs-comment"># 创建后处理类的实例</span>
postprocessor = PostprocessYOLO(**postprocessor_args)
<span class="hljs-comment"># 运行后处理算法，并得到检测到对象的bounding box</span>
boxes, classes, scores = postprocessor.process(trt_outputs, (shape_orig_WH))

<span class="hljs-string">''' 5 - 在原始输入图像上将检测框标记，并保存png文件 '''</span>
obj_detected_img = draw_bboxes(image_raw, boxes, scores, classes, ALL_CATEGORIES)
output_image_path = <span class="hljs-string">'dog_bboxes.png'</span>
obj_detected_img.save(output_image_path, <span class="hljs-string">'PNG'</span>)

print(<span class="hljs-string">'Saved image with bounding boxes of detected objects to {}.'</span>.format(output_image_path))

if name == ‘main’:
main()

运行程序及结果：

python onnx_to_tensorrt.py

此时文件目录为：

.
├── coco_labels.txt
├── data_processing.py
├── dog_bboxes.png
├── dog.jpg
├── onnx_to_tensorrt.py
├── __pycache__
│   ├── data_processing.cpython-35.pyc
│   └── yolov3_to_onnx.cpython-35.pyc
├── README.md
├── requirements.txt
├── yolov3.cfg
├── yolov3.onnx
├── yolov3_to_onnx.py
├── yolov3.trt
└── yolov3.weights

最后我们来看下data_processing.py

import math
from PIL import Image
import numpy as np

# YOLOv3-608 has been trained with these 80 categories from COCO:
# Lin, Tsung-Yi, et al. “Microsoft COCO: Common Objects in Context.”
# European Conference on Computer Vision. Springer, Cham, 2014.

def load_label_categories(label_file_path):
categories = [line.rstrip(’\n’) for line in open(label_file_path)]
return categories

LABEL_FILE_PATH = ‘coco_labels.txt’
ALL_CATEGORIES = load_label_categories(LABEL_FILE_PATH)

# 确定有80个类别:
CATEGORY_NUM = len(ALL_CATEGORIES)
assert CATEGORY_NUM == 80

class PreprocessYOLO(object):
“”“装载图像，然后reshape成yolov3-608需要的分辨率.
“””

<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">__init__</span><span class="hljs-params">(self, yolo_input_resolution)</span>:</span>
    <span class="hljs-string">"""指定yolov3的输入分辨率.

    Keyword arguments:
    yolo_input_resolution -- two-dimensional tuple with the target network's (spatial)
    input resolution in HW order
    """</span>
    self.yolo_input_resolution = yolo_input_resolution

<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">process</span><span class="hljs-params">(self, input_image_path)</span>:</span>
    <span class="hljs-string">"""载入图像，然后进行预处理，如resize，归一化等等

    Keyword arguments:
    input_image_path -- string path of the image to be loaded
    """</span>
    image_raw, image_resized = self._load_and_resize(input_image_path)
    image_preprocessed = self._shuffle_and_normalize(image_resized)
    <span class="hljs-keyword">return</span> image_raw, image_preprocessed

<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">_load_and_resize</span><span class="hljs-params">(self, input_image_path)</span>:</span>
    <span class="hljs-string">"""对图像进行resize，然后返回numpy对象

    Keyword arguments:
    input_image_path -- string path of the image to be loaded
    """</span>

    image_raw = Image.open(input_image_path)

    new_resolution = (
        self.yolo_input_resolution[<span class="hljs-number">1</span>],
        self.yolo_input_resolution[<span class="hljs-number">0</span>])
    image_resized = image_raw.resize(
        new_resolution, resample=Image.BICUBIC)
    image_resized = np.array(image_resized, dtype=np.float32, order=<span class="hljs-string">'C'</span>)
    <span class="hljs-keyword">return</span> image_raw, image_resized

<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">_shuffle_and_normalize</span><span class="hljs-params">(self, image)</span>:</span>
    <span class="hljs-string">"""将图像归一化到[0,1]之间，然后将HWC结构转换成NCHW结构

    Keyword arguments:
    image -- image as three-dimensional NumPy float array, in HWC format
    """</span>
    image /= <span class="hljs-number">255.0</span>
    <span class="hljs-comment"># HWC -&gt; CHW :</span>
    image = np.transpose(image, [<span class="hljs-number">2</span>, <span class="hljs-number">0</span>, <span class="hljs-number">1</span>])
    <span class="hljs-comment"># CHW -&gt; NCHW </span>
    image = np.expand_dims(image, axis=<span class="hljs-number">0</span>)
    <span class="hljs-comment"># j将图像转换成row-major order,如 "C order":</span>
    image = np.array(image, dtype=np.float32, order=<span class="hljs-string">'C'</span>)
    <span class="hljs-keyword">return</span> image

class PostprocessYOLO(object):
“”“后处理yolov3-608的三个输出tensor.”""

<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">__init__</span><span class="hljs-params">(self,
             yolo_masks,
             yolo_anchors,
             obj_threshold,
             nms_threshold,
             yolo_input_resolution)</span>:</span>
    <span class="hljs-string">"""Initialize with all values that will be kept when processing several frames.
    Assuming 3 outputs of the network in the case of (large) YOLOv3.

    Keyword arguments:
    yolo_masks -- a list of 3 three-dimensional tuples for the YOLO masks
    yolo_anchors -- a list of 9 two-dimensional tuples for the YOLO anchors
    object_threshold -- threshold for object coverage, float value between 0 and 1
    nms_threshold -- threshold for non-max suppression algorithm,
    float value between 0 and 1
    input_resolution_yolo -- two-dimensional tuple with the target network's (spatial)
    input resolution in HW order
    """</span>
    self.masks = yolo_masks
    self.anchors = yolo_anchors
    self.object_threshold = obj_threshold
    self.nms_threshold = nms_threshold
    self.input_resolution_yolo = yolo_input_resolution

<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">process</span><span class="hljs-params">(self, outputs, resolution_raw)</span>:</span>
    <span class="hljs-string">"""Take the YOLOv3 outputs generated from a TensorRT forward pass, post-process them
    and return a list of bounding boxes for detected object together with their category
    and their confidences in separate lists.

    Keyword arguments:
    outputs -- outputs from a TensorRT engine in NCHW format
    resolution_raw -- the original spatial resolution from the input PIL image in WH order
    """</span>
    outputs_reshaped = list()
    <span class="hljs-keyword">for</span> output <span class="hljs-keyword">in</span> outputs:
        outputs_reshaped.append(self._reshape_output(output))

    boxes, categories, confidences = self._process_yolo_output(
        outputs_reshaped, resolution_raw)

    <span class="hljs-keyword">return</span> boxes, categories, confidences

<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">_reshape_output</span><span class="hljs-params">(self, output)</span>:</span>
    <span class="hljs-string">"""Reshape a TensorRT output from NCHW to NHWC format (with expected C=255),
    and then return it in (height,width,3,85) dimensionality after further reshaping.

    Keyword argument:
    output -- an output from a TensorRT engine after inference
    """</span>
    output = np.transpose(output, [<span class="hljs-number">0</span>, <span class="hljs-number">2</span>, <span class="hljs-number">3</span>, <span class="hljs-number">1</span>])
    _, height, width, _ = output.shape
    dim1, dim2 = height, width
    dim3 = <span class="hljs-number">3</span>
    <span class="hljs-comment"># There are CATEGORY_NUM=80 object categories:</span>
    dim4 = (<span class="hljs-number">4</span> + <span class="hljs-number">1</span> + CATEGORY_NUM)
    <span class="hljs-keyword">return</span> np.reshape(output, (dim1, dim2, dim3, dim4))

<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">_process_yolo_output</span><span class="hljs-params">(self, outputs_reshaped, resolution_raw)</span>:</span>
    <span class="hljs-string">"""Take in a list of three reshaped YOLO outputs in (height,width,3,85) shape and return
    return a list of bounding boxes for detected object together with their category and their
    confidences in separate lists.

    Keyword arguments:
    outputs_reshaped -- list of three reshaped YOLO outputs as NumPy arrays
    with shape (height,width,3,85)
    resolution_raw -- the original spatial resolution from the input PIL image in WH order
    """</span>

    <span class="hljs-comment"># E.g. in YOLOv3-608, there are three output tensors, which we associate with their</span>
    <span class="hljs-comment"># respective masks. Then we iterate through all output-mask pairs and generate candidates</span>
    <span class="hljs-comment"># for bounding boxes, their corresponding category predictions and their confidences:</span>
    boxes, categories, confidences = list(), list(), list()
    <span class="hljs-keyword">for</span> output, mask <span class="hljs-keyword">in</span> zip(outputs_reshaped, self.masks):
        box, category, confidence = self._process_feats(output, mask)
        box, category, confidence = self._filter_boxes(box, category, confidence)
        boxes.append(box)
        categories.append(category)
        confidences.append(confidence)

    boxes = np.concatenate(boxes)
    categories = np.concatenate(categories)
    confidences = np.concatenate(confidences)

    <span class="hljs-comment"># Scale boxes back to original image shape:</span>
    width, height = resolution_raw
    image_dims = [width, height, width, height]
    boxes = boxes * image_dims

    <span class="hljs-comment"># Using the candidates from the previous (loop) step, we apply the non-max suppression</span>
    <span class="hljs-comment"># algorithm that clusters adjacent bounding boxes to a single bounding box:</span>
    nms_boxes, nms_categories, nscores = list(), list(), list()
    <span class="hljs-keyword">for</span> category <span class="hljs-keyword">in</span> set(categories):
        idxs = np.where(categories == category)
        box = boxes[idxs]
        category = categories[idxs]
        confidence = confidences[idxs]

        keep = self._nms_boxes(box, confidence)

        nms_boxes.append(box[keep])
        nms_categories.append(category[keep])
        nscores.append(confidence[keep])

    <span class="hljs-keyword">if</span> <span class="hljs-keyword">not</span> nms_categories <span class="hljs-keyword">and</span> <span class="hljs-keyword">not</span> nscores:
        <span class="hljs-keyword">return</span> <span class="hljs-keyword">None</span>, <span class="hljs-keyword">None</span>, <span class="hljs-keyword">None</span>

    boxes = np.concatenate(nms_boxes)
    categories = np.concatenate(nms_categories)
    confidences = np.concatenate(nscores)

    <span class="hljs-keyword">return</span> boxes, categories, confidences

<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">_process_feats</span><span class="hljs-params">(self, output_reshaped, mask)</span>:</span>
    <span class="hljs-string">"""Take in a reshaped YOLO output in height,width,3,85 format together with its
    corresponding YOLO mask and return the detected bounding boxes, the confidence,
    and the class probability in each cell/pixel.

    Keyword arguments:
    output_reshaped -- reshaped YOLO output as NumPy arrays with shape (height,width,3,85)
    mask -- 2-dimensional tuple with mask specification for this output
    """</span>

    <span class="hljs-comment"># Two in-line functions required for calculating the bounding box</span>
    <span class="hljs-comment"># descriptors:</span>
    <span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">sigmoid</span><span class="hljs-params">(value)</span>:</span>
        <span class="hljs-string">"""Return the sigmoid of the input."""</span>
        <span class="hljs-keyword">return</span> <span class="hljs-number">1.0</span> / (<span class="hljs-number">1.0</span> + math.exp(-value))

    <span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">exponential</span><span class="hljs-params">(value)</span>:</span>
        <span class="hljs-string">"""Return the exponential of the input."""</span>
        <span class="hljs-keyword">return</span> math.exp(value)

    <span class="hljs-comment"># Vectorized calculation of above two functions:</span>
    sigmoid_v = np.vectorize(sigmoid)
    exponential_v = np.vectorize(exponential)

    grid_h, grid_w, _, _ = output_reshaped.shape

    anchors = [self.anchors[i] <span class="hljs-keyword">for</span> i <span class="hljs-keyword">in</span> mask]

    <span class="hljs-comment"># Reshape to N, height, width, num_anchors, box_params:</span>
    anchors_tensor = np.reshape(anchors, [<span class="hljs-number">1</span>, <span class="hljs-number">1</span>, len(anchors), <span class="hljs-number">2</span>])
    box_xy = sigmoid_v(output_reshaped[..., :<span class="hljs-number">2</span>])
    box_wh = exponential_v(output_reshaped[..., <span class="hljs-number">2</span>:<span class="hljs-number">4</span>]) * anchors_tensor
    box_confidence = sigmoid_v(output_reshaped[..., <span class="hljs-number">4</span>])

    box_confidence = np.expand_dims(box_confidence, axis=<span class="hljs-number">-1</span>)
    box_class_probs = sigmoid_v(output_reshaped[..., <span class="hljs-number">5</span>:])

    col = np.tile(np.arange(<span class="hljs-number">0</span>, grid_w), grid_w).reshape(<span class="hljs-number">-1</span>, grid_w)
    row = np.tile(np.arange(<span class="hljs-number">0</span>, grid_h).reshape(<span class="hljs-number">-1</span>, <span class="hljs-number">1</span>), grid_h)

    col = col.reshape(grid_h, grid_w, <span class="hljs-number">1</span>, <span class="hljs-number">1</span>).repeat(<span class="hljs-number">3</span>, axis=<span class="hljs-number">-2</span>)
    row = row.reshape(grid_h, grid_w, <span class="hljs-number">1</span>, <span class="hljs-number">1</span>).repeat(<span class="hljs-number">3</span>, axis=<span class="hljs-number">-2</span>)
    grid = np.concatenate((col, row), axis=<span class="hljs-number">-1</span>)

    box_xy += grid
    box_xy /= (grid_w, grid_h)
    box_wh /= self.input_resolution_yolo
    box_xy -= (box_wh / <span class="hljs-number">2.</span>)
    boxes = np.concatenate((box_xy, box_wh), axis=<span class="hljs-number">-1</span>)

    <span class="hljs-comment"># boxes: centroids, box_confidence: confidence level, box_class_probs:</span>
    <span class="hljs-comment"># class confidence</span>
    <span class="hljs-keyword">return</span> boxes, box_confidence, box_class_probs

<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">_filter_boxes</span><span class="hljs-params">(self, boxes, box_confidences, box_class_probs)</span>:</span>
    <span class="hljs-string">"""Take in the unfiltered bounding box descriptors and discard each cell
    whose score is lower than the object threshold set during class initialization.

    Keyword arguments:
    boxes -- bounding box coordinates with shape (height,width,3,4); 4 for
    x,y,height,width coordinates of the boxes
    box_confidences -- bounding box confidences with shape (height,width,3,1); 1 for as
    confidence scalar per element
    box_class_probs -- class probabilities with shape (height,width,3,CATEGORY_NUM)

    """</span>
    box_scores = box_confidences * box_class_probs
    box_classes = np.argmax(box_scores, axis=<span class="hljs-number">-1</span>)
    box_class_scores = np.max(box_scores, axis=<span class="hljs-number">-1</span>)
    pos = np.where(box_class_scores &gt;= self.object_threshold)

    boxes = boxes[pos]
    classes = box_classes[pos]
    scores = box_class_scores[pos]

    <span class="hljs-keyword">return</span> boxes, classes, scores

<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">_nms_boxes</span><span class="hljs-params">(self, boxes, box_confidences)</span>:</span>
    <span class="hljs-string">"""Apply the Non-Maximum Suppression (NMS) algorithm on the bounding boxes with their
    confidence scores and return an array with the indexes of the bounding boxes we want to
    keep (and display later).

    Keyword arguments:
    boxes -- a NumPy array containing N bounding-box coordinates that survived filtering,
    with shape (N,4); 4 for x,y,height,width coordinates of the boxes
    box_confidences -- a Numpy array containing the corresponding confidences with shape N
    """</span>
    x_coord = boxes[:, <span class="hljs-number">0</span>]
    y_coord = boxes[:, <span class="hljs-number">1</span>]
    width = boxes[:, <span class="hljs-number">2</span>]
    height = boxes[:, <span class="hljs-number">3</span>]

    areas = width * height
    ordered = box_confidences.argsort()[::<span class="hljs-number">-1</span>]

    keep = list()
    <span class="hljs-keyword">while</span> ordered.size &gt; <span class="hljs-number">0</span>:
        <span class="hljs-comment"># Index of the current element:</span>
        i = ordered[<span class="hljs-number">0</span>]
        keep.append(i)
        xx1 = np.maximum(x_coord[i], x_coord[ordered[<span class="hljs-number">1</span>:]])
        yy1 = np.maximum(y_coord[i], y_coord[ordered[<span class="hljs-number">1</span>:]])
        xx2 = np.minimum(x_coord[i] + width[i], x_coord[ordered[<span class="hljs-number">1</span>:]] + width[ordered[<span class="hljs-number">1</span>:]])
        yy2 = np.minimum(y_coord[i] + height[i], y_coord[ordered[<span class="hljs-number">1</span>:]] + height[ordered[<span class="hljs-number">1</span>:]])

        width1 = np.maximum(<span class="hljs-number">0.0</span>, xx2 - xx1 + <span class="hljs-number">1</span>)
        height1 = np.maximum(<span class="hljs-number">0.0</span>, yy2 - yy1 + <span class="hljs-number">1</span>)
        intersection = width1 * height1
        union = (areas[i] + areas[ordered[<span class="hljs-number">1</span>:]] - intersection)

        <span class="hljs-comment"># Compute the Intersection over Union (IoU) score:</span>
        iou = intersection / union

        <span class="hljs-comment"># The goal of the NMS algorithm is to reduce the number of adjacent bounding-box</span>
        <span class="hljs-comment"># candidates to a minimum. In this step, we keep only those elements whose overlap</span>
        <span class="hljs-comment"># with the current bounding box is lower than the threshold:</span>
        indexes = np.where(iou &lt;= self.nms_threshold)[<span class="hljs-number">0</span>]
        ordered = ordered[indexes + <span class="hljs-number">1</span>]

    keep = np.array(keep)
    <span class="hljs-keyword">return</span> keep