yolov3_tiny.onnx转trt采用tensorrt加速模型推理
既然上一篇博客都把yolov3-tiny.weights转onnx做了，推理也测了。那么呢，就再直接转个trt模型吧。这样感觉博客的内容就更加连贯了吧，实用性貌似会更加强吧。
（如果没看过yolov3-tiny转onnx这篇博客的，请点这，带你飞过去）
这篇博客的内容是接着上一篇博客写的，所以呢，这里就直接进入主题，上代码！！
本文的目录：
- onnx模型转trt文件（yolov3-tiny.onnx）；
- 实用tensorrt进行推理；
- 测试推理效果（视频文件）；
首先呢，你得要安装好tensorrt，至于怎么安装很多博客都有介绍，在本篇博客中使用的tensorrt版本为：
Python 3.5.6 |Anaconda, Inc.| (default, Aug 26 2018, 21:41:56)
[GCC 7.3.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import tensorrt as trt
>>> print(trt.__version__)
6.0.1.5
测试的硬件配置：
GTX2060（laptop）-6G
i7-9750H
16G-DDR4 2666MHz
一. onnx模型转trt文件
直接上代码，运行就完事（代码均可将yolov3.onnx、yolov3-tiny.onnx转为trt文件）
onnx_to_trt.py
# -*-coding: utf-8-*-
# author: HXY
# 2019-12-24
"""
tensorrt6.0
"""
import os
import tensorrt as trt
TRT_LOGGER = trt.Logger()
class GetTrt(object):
def __init__(self, onnx_file_path, trt_save_path):
self.onnx_file_path = onnx_file_path
self.batch_size = 1
self.fp16_on = True
self.trt_save_path = trt_save_path
def build_engine(self):
with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.OnnxParser(network,
TRT_LOGGER) as parser:
builder.max_workspace_size = 1 << 30 # 30:1GB; 28:256MiB
builder.max_batch_size = 1
builder.fp16_mode = self.fp16_on
if not os.path.exists(self.onnx_file_path):
print("Onnx file not found")
exit(0)
print("loading onnx file from path {}".format(self.onnx_file_path))
with open(self.onnx_file_path, 'rb') as model:
parser.parse(model.read())
print("Completed parsing of onnx file....")
print("building an engine..this may take a while....")
# network.get_input(0).shape = [1, 3, 416, 416]
engine = builder.build_cuda_engine(network)
with open(self.trt_save_path, 'wb') as f:
f.write(engine.serialize())
print("create engine completed")
"""
test function
"""
if __name__ == '__main__':
test = GetTrt('./yolov3-tiny.onnx',
'./yolov3-tiny.trt')
test.build_engine()
你运行成功后，你就可以得到一份trt文件，这个时候你就可以开始进行推理测试啦。
这里贴一张运行成功的截图（额，有个警告，请忽略 哈哈）：
二.使用tensorrt进行推理
既然进行推理，那还是需要一点代码的，这些代码建议放在一个文件夹内，供后面推理时调用；
首先是common.py ,这些代码，用就完事了，喜欢的朋友可以自行研究研究！
#
# Copyright 1993-2019 NVIDIA Corporation. All rights reserved.
#
# NOTICE TO LICENSEE:
#
# This source code and/or documentation ("Licensed Deliverables") are
# subject to NVIDIA intellectual property rights under U.S. and
# international Copyright laws.
#
# These Licensed Deliverables contained herein is PROPRIETARY and
# CONFIDENTIAL to NVIDIA and is being provided under the terms and
# conditions of a form of NVIDIA software license agreement by and
# between NVIDIA and Licensee ("License Agreement") or electronically
# accepted by Licensee. Notwithstanding any terms or conditions to
# the contrary in the License Agreement, reproduction or disclosure
# of the Licensed Deliverables to any third party without the express
# written consent of NVIDIA is prohibited.
#
# NOTWITHSTANDING ANY TERMS OR CONDITIONS TO THE CONTRARY IN THE
# LICENSE AGREEMENT, NVIDIA MAKES NO REPRESENTATION ABOUT THE
# SUITABILITY OF THESE LICENSED DELIVERABLES FOR ANY PURPOSE. IT IS
# PROVIDED "AS IS" WITHOUT EXPRESS OR IMPLIED WARRANTY OF ANY KIND.
# NVIDIA DISCLAIMS ALL WARRANTIES WITH REGARD TO THESE LICENSED
# DELIVERABLES, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY,
# NONINFRINGEMENT, AND FITNESS FOR A PARTICULAR PURPOSE.
# NOTWITHSTANDING ANY TERMS OR CONDITIONS TO THE CONTRARY IN THE
# LICENSE AGREEMENT, IN NO EVENT SHALL NVIDIA BE LIABLE FOR ANY
# SPECIAL, INDIRECT, INCIDENTAL, OR CONSEQUENTIAL DAMAGES, OR ANY
# DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS,
# WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS
# ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE
# OF THESE LICENSED DELIVERABLES.
#
# U.S. Government End Users. These Licensed Deliverables are a
# "commercial item" as that term is defined at 48 C.F.R. 2.101 (OCT
# 1995), consisting of "commercial computer software" and "commercial
# computer software documentation" as such terms are used in 48
# C.F.R. 12.212 (SEPT 1995) and is provided to the U.S. Government
# only as a commercial end item. Consistent with 48 C.F.R.12.212 and
# 48 C.F.R. 227.7202-1 through 227.7202-4 (JUNE 1995), all
# U.S. Government End Users acquire the Licensed Deliverables with
# only those rights set forth herein.
#
# Any use of the Licensed Deliverables in individual and commercial
# software must include, in the user documentation and internal
# comments to the code, the above Disclaimer and U.S. Government End
# Users Notice.
#
from itertools import chain
import argparse
import os
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
import tensorrt as trt
try:
# Sometimes python2 does not understand FileNotFoundError
FileNotFoundError
except NameError:
FileNotFoundError = IOError
def GiB(val):
return val * 1 << 30
def find_sample_data(description="Runs a TensorRT Python sample", subfolder="", find_files=[]):
'''
Parses sample arguments.
Args:
description (str): Description of the sample.
subfolder (str): The subfolder containing data relevant to this sample
find_files (str): A list of filenames to find. Each filename will be replaced with an absolute path.
Returns:
str: Path of data directory.
'''
# Standard command-line arguments for all samples.
kDEFAULT_DATA_ROOT = os.path.join(os.sep, "usr", "src", "tensorrt", "data")
parser = argparse.ArgumentParser(description=description, formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument("-d", "--datadir",
help="Location of the TensorRT sample data directory, and any additional data directories.",
action="append", default=[kDEFAULT_DATA_ROOT])
args, _ = parser.parse_known_args()
def get_data_path(data_dir):
# If the subfolder exists, append it to the path, otherwise use the provided path as-is.
data_path = os.path.join(data_dir, subfolder)
if not os.path.exists(data_path):
print("WARNING: " + data_path + " does not exist. Trying " + data_dir + " instead.")
data_path = data_dir
# Make sure data directory exists.
if not (os.path.exists(data_path)):
print("WARNING: {:} does not exist. Please provide the correct data path with the -d option.".format(
data_path))
return data_path
data_paths = [get_data_path(data_dir) for data_dir in args.datadir]
return data_paths, locate_files(data_paths, find_files)
def locate_files(data_paths, filenames):
"""
Locates the specified files in the specified data directories.
If a file exists in multiple data directories, the first directory is used.
Args:
data_paths (List[str]): The data directories.
filename (List[str]): The names of the files to find.
Returns:
List[str]: The absolute paths of the files.
Raises:
FileNotFoundError if a file could not be located.
"""
found_files = [None] * len(filenames)
for data_path in data_paths:
# Find all requested files.
for index, (found, filename) in enumerate(zip(found_files, filenames)):
if not found:
file_path = os.path.abspath(os.path.join(data_path, filename))
if os.path.exists(file_path):
found_files[index] = file_path
# Check that all files were found
for f, filename in zip(found_files, filenames):
if not f or not os.path.exists(f):
raise FileNotFoundError("Could not find {:}. Searched in data paths: {:}".format(filename, data_paths))
return found_files
# Simple helper data class that's a little nicer to use than a 2-tuple.
class HostDeviceMem(object):
def __init__(self, host_mem, device_mem):
self.host = host_mem
self.device = device_mem
def __str__(self):
return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device)
def __repr__(self):
return self.__str__()
# Allocates all buffers required for an engine, i.e. host/device inputs/outputs.
def allocate_buffers(engine):
inputs = []
outputs = []
bindings = []
stream = cuda.Stream()
for binding in engine:
size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
dtype = trt.nptype(engine.get_binding_dtype(binding))
# Allocate host and device buffers
host_mem = cuda.pagelocked_empty(size, dtype)
device_mem = cuda.mem_alloc(host_mem.nbytes)
# Append the device buffer to device bindings.
bindings.append(int(device_mem))
# Append to the appropriate list.
if engine.binding_is_input(binding):
inputs.append(HostDeviceMem(host_mem, device_mem))
else:
outputs.append(HostDeviceMem(host_mem, device_mem))
return inputs, outputs, bindings, stream
# This function is generalized for multiple inputs/outputs.
# inputs and outputs are expected to be lists of HostDeviceMem objects.
def do_inference(context, bindings, inputs, outputs, stream, batch_size=1):
# Transfer input data to the GPU.
[cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
# Run inference.
context.execute_async(batch_size=batch_size, bindings=bindings, stream_handle=stream.handle)
# Transfer predictions back from the GPU.
[cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
# Synchronize the stream
stream.synchronize()
# Return only the host outputs.
return [out.host for out in outputs]
数据预处理部分代码：
data_processing.py
#
# Copyright 1993-2019 NVIDIA Corporation. All rights reserved.
#
# NOTICE TO LICENSEE:
#
# This source code and/or documentation ("Licensed Deliverables") are
# subject to NVIDIA intellectual property rights under U.S. and
# international Copyright laws.
#
# These Licensed Deliverables contained herein is PROPRIETARY and
# CONFIDENTIAL to NVIDIA and is being provided under the terms and
# conditions of a form of NVIDIA software license agreement by and
# between NVIDIA and Licensee ("License Agreement") or electronically
# accepted by Licensee. Notwithstanding any terms or conditions to
# the contrary in the License Agreement, reproduction or disclosure
# of the Licensed Deliverables to any third party without the express
# written consent of NVIDIA is prohibited.
#
# NOTWITHSTANDING ANY TERMS OR CONDITIONS TO THE CONTRARY IN THE
# LICENSE AGREEMENT, NVIDIA MAKES NO REPRESENTATION ABOUT THE
# SUITABILITY OF THESE LICENSED DELIVERABLES FOR ANY PURPOSE. IT IS
# PROVIDED "AS IS" WITHOUT EXPRESS OR IMPLIED WARRANTY OF ANY KIND.
# NVIDIA DISCLAIMS ALL WARRANTIES WITH REGARD TO THESE LICENSED
# DELIVERABLES, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY,
# NONINFRINGEMENT, AND FITNESS FOR A PARTICULAR PURPOSE.
# NOTWITHSTANDING ANY TERMS OR CONDITIONS TO THE CONTRARY IN THE
# LICENSE AGREEMENT, IN NO EVENT SHALL NVIDIA BE LIABLE FOR ANY
# SPECIAL, INDIRECT, INCIDENTAL, OR CONSEQUENTIAL DAMAGES, OR ANY
# DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS,
# WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS
# ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE
# OF THESE LICENSED DELIVERABLES.
#
# U.S. Government End Users. These Licensed Deliverables are a
# "commercial item" as that term is defined at 48 C.F.R. 2.101 (OCT
# 1995), consisting of "commercial computer software" and "commercial
# computer software documentation" as such terms are used in 48
# C.F.R. 12.212 (SEPT 1995) and is provided to the U.S. Government
# only as a commercial end item. Consistent with 48 C.F.R.12.212 and
# 48 C.F.R. 227.7202-1 through 227.7202-4 (JUNE 1995), all
# U.S. Government End Users acquire the Licensed Deliverables with
# only those rights set forth herein.
#
# Any use of the Licensed Deliverables in individual and commercial
# software must include, in the user documentation and internal
# comments to the code, the above Disclaimer and U.S. Government End
# Users Notice.
#
# -*-coding: utf-8-*-
import math
from PIL import Image
import numpy as np
import os
# 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 = os.path.join(os.path.dirname(os.path.realpath(__file__)), '../config/lxz.txt')
ALL_CATEGORIES = load_label_categories(LABEL_FILE_PATH)
# Let's make sure that there are 80 classes, as expected for the COCO data set:
CATEGORY_NUM = len(ALL_CATEGORIES)
assert CATEGORY_NUM == 3
# 增加视频处理函数
class VideoPreprocessYOLO(object):
"""A simple class for loading images with PIL and reshaping them to the specified
input resolution for YOLOv3-608.
"""
def __init__(self, yolo_input_resolution):
"""Initialize with the input resolution for YOLOv3, which will stay fixed in this sample.
Keyword arguments:
yolo_input_resolution -- two-dimensional tuple with the target network's (spatial)
input resolution in HW order
"""
self.yolo_input_resolution = yolo_input_resolution
def process(self, input_image_raw):
"""Load an image from the specified input path,
and return it together with a pre-processed version required for feeding it into a
YOLOv3 network.
Keyword arguments:
input_image_path -- string path of the image to be loaded
"""
image_raw, image_resized = self._load_and_resize(input_image_raw)
image_preprocessed = self._shuffle_and_normalize(image_resized)
return image_raw, image_preprocessed
def _load_and_resize(self, input_image_raw):
"""Load an image from the specified path and resize it to the input resolution.
Return the input image before resizing as a PIL Image (required for visualization),
and the resized image as a NumPy float array.
Keyword arguments:
input_image_path -- string path of the image to be loaded
"""
image_raw = input_image_raw
# Expecting yolo_input_resolution in (height, width) format, adjusting to PIL
# convention (width, height) in PIL:
new_resolution = (
self.yolo_input_resolution[1],
self.yolo_input_resolution[0])
image_resized = image_raw.resize(
new_resolution, resample=Image.BICUBIC)
image_resized = np.array(image_resized, dtype=np.float32, order='C')
return image_raw, image_resized
def _shuffle_and_normalize(self, image):
"""Normalize a NumPy array representing an image to the range [0, 1], and
convert it from HWC format ("channels last") to NCHW format ("channels first"
with leading batch dimension).
Keyword arguments:
image -- image as three-dimensional NumPy float array, in HWC format
"""
image /= 255.0
# HWC to CHW format:
image = np.transpose(image, [2, 0, 1])
# CHW to NCHW format
image = np.expand_dims(image, axis=0)
# Convert the image to row-major order, also known as "C order":
image = np.array(image, dtype=np.float32, order='C')
return image
class PreprocessYOLO(object):
"""A simple class for loading images with PIL and reshaping them to the specified
input resolution for YOLOv3-608.
"""
def __init__(self, yolo_input_resolution):
"""Initialize with the input resolution for YOLOv3, which will stay fixed in this sample.
Keyword arguments:
yolo_input_resolution -- two-dimensional tuple with the target network's (spatial)
input resolution in HW order
"""
self.yolo_input_resolution = yolo_input_resolution
def process(self, input_image_path):
"""Load an image from the specified input path,
and return it together with a pre-processed version required for feeding it into a
YOLOv3 network.
Keyword arguments:
input_image_path -- string path of the image to be loaded
"""
image_raw, image_resized = self._load_and_resize(input_image_path)
image_preprocessed = self._shuffle_and_normalize(image_resized)
return image_raw, image_preprocessed
def _load_and_resize(self, input_image_path):
"""Load an image from the specified path and resize it to the input resolution.
Return the input image before resizing as a PIL Image (required for visualization),
and the resized image as a NumPy float array.
Keyword arguments:
input_image_path -- string path of the image to be loaded
"""
image_raw = Image.open(input_image_path)
# Expecting yolo_input_resolution in (height, width) format, adjusting to PIL
# convention (width, height) in PIL:
new_resolution = (
self.yolo_input_resolution[1],
self.yolo_input_resolution[0])
image_resized = image_raw.resize(
new_resolution, resample=Image.BICUBIC)
image_resized = np.array(image_resized, dtype=np.float32, order='C')
return image_raw, image_resized
def _shuffle_and_normalize(self, image):
"""Normalize a NumPy array representing an image to the range [0, 1], and
convert it from HWC format ("channels last") to NCHW format ("channels first"
with leading batch dimension).
Keyword arguments:
image -- image as three-dimensional NumPy float array, in HWC format
"""
image /= 255.0
# HWC to CHW format:
image = np.transpose(image, [2, 0, 1])
# CHW to NCHW format
image = np.expand_dims(image, axis=0)
# Convert the image to row-major order, also known as "C order":
image = np.array(image, dtype=np.float32, order='C')
return image
class PostprocessYOLO(object):
"""Class for post-processing the three outputs tensors from YOLOv3-608."""
def __init__(self,
yolo_masks,
yolo_anchors,
obj_threshold,
nms_threshold,
yolo_input_resolution):
"""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
"""
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
def process(self, outputs, resolution_raw):
"""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
"""
outputs_reshaped = list()
for output in outputs:
outputs_reshaped.append(self._reshape_output(output))
boxes, categories, confidences = self._process_yolo_output(
outputs_reshaped, resolution_raw)
return boxes, categories, confidences
def _reshape_output(self, output):
"""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
"""
output = np.transpose(output, [0, 2, 3, 1])
_, height, width, _ = output.shape
dim1, dim2 = height, width
dim3 = 3
# There are CATEGORY_NUM=80 object categories:
dim4 = (4 + 1 + CATEGORY_NUM)
return np.reshape(output, (dim1, dim2, dim3, dim4))
def _process_yolo_output(self, outputs_reshaped, resolution_raw):
"""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
"""
# E.g. in YOLOv3-608, there are three output tensors, which we associate with their
# respective masks. Then we iterate through all output-mask pairs and generate candidates
# for bounding boxes, their corresponding category predictions and their confidences:
boxes, categories, confidences = list(), list(), list()
for output, mask in 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)
# Scale boxes back to original image shape:
width, height = resolution_raw
image_dims = [width, height, width, height]
boxes = boxes * image_dims
# Using the candidates from the previous (loop) step, we apply the non-max suppression
# algorithm that clusters adjacent bounding boxes to a single bounding box:
nms_boxes, nms_categories, nscores = list(), list(), list()
for category in 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])
if not nms_categories and not nscores:
return None, None, None
boxes = np.concatenate(nms_boxes)
categories = np.concatenate(nms_categories)
confidences = np.concatenate(nscores)
return boxes, categories, confidences
def _process_feats(self, output_reshaped, mask):
"""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
"""
# Two in-line functions required for calculating the bounding box
# descriptors:
def sigmoid(value):
"""Return the sigmoid of the input."""
return 1.0 / (1.0 + math.exp(-value))
def exponential(value):
"""Return the exponential of the input."""
return math.exp(value)
# Vectorized calculation of above two functions:
sigmoid_v = np.vectorize(sigmoid)
exponential_v = np.vectorize(exponential)
grid_h, grid_w, _, _ = output_reshaped.shape
anchors = [self.anchors[i] for i in mask]
# Reshape to N, height, width, num_anchors, box_params:
anchors_tensor = np.reshape(anchors, [1, 1, len(anchors), 2])
box_xy = sigmoid_v(output_reshaped[..., :2])
box_wh = exponential_v(output_reshaped[..., 2:4]) * anchors_tensor
box_confidence = sigmoid_v(output_reshaped[..., 4])
box_confidence = np.expand_dims(box_confidence, axis=-1)
box_class_probs = sigmoid_v(output_reshaped[..., 5:])
col = np.tile(np.arange(0, grid_w), grid_w).reshape(-1, grid_w)
row = np.tile(np.arange(0, grid_h).reshape(-1, 1), grid_h)
col = col.reshape(grid_h, grid_w, 1, 1).repeat(3, axis=-2)
row = row.reshape(grid_h, grid_w, 1, 1).repeat(3, axis=-2)
grid = np.concatenate((col, row), axis=-1)
box_xy += grid
box_xy /= (grid_w, grid_h)
box_wh /= self.input_resolution_yolo
box_xy -= (box_wh / 2.)
boxes = np.concatenate((box_xy, box_wh), axis=-1)
# boxes: centroids, box_confidence: confidence level, box_class_probs:
# class confidence
return boxes, box_confidence, box_class_probs
def _filter_boxes(self, boxes, box_confidences, box_class_probs):
"""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)
"""
box_scores = box_confidences * box_class_probs
box_classes = np.argmax(box_scores, axis=-1)
box_class_scores = np.max(box_scores, axis=-1)
pos = np.where(box_class_scores >= self.object_threshold)
boxes = boxes[pos]
classes = box_classes[pos]
scores = box_class_scores[pos]
return boxes, classes, scores
def _nms_boxes(self, boxes, box_confidences):
"""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
"""
x_coord = boxes[:, 0]
y_coord = boxes[:, 1]
width = boxes[:, 2]
height = boxes[:, 3]
areas = width * height
ordered = box_confidences.argsort()[::-1]
keep = list()
while ordered.size > 0:
# Index of the current element:
i = ordered[0]
keep.append(i)
xx1 = np.maximum(x_coord[i], x_coord[ordered[1:]])
yy1 = np.maximum(y_coord[i], y_coord[ordered[1:]])
xx2 = np.minimum(x_coord[i] + width[i], x_coord[ordered[1:]] + width[ordered[1:]])
yy2 = np.minimum(y_coord[i] + height[i], y_coord[ordered[1:]] + height[ordered[1:]])
width1 = np.maximum(0.0, xx2 - xx1 + 1)
height1 = np.maximum(0.0, yy2 - yy1 + 1)
intersection = width1 * height1
union = (areas[i] + areas[ordered[1:]] - intersection)
# Compute the Intersection over Union (IoU) score:
iou = intersection / union
# The goal of the NMS algorithm is to reduce the number of adjacent bounding-box
# candidates to a minimum. In this step, we keep only those elements whose overlap
# with the current bounding box is lower than the threshold:
indexes = np.where(iou <= self.nms_threshold)[0]
ordered = ordered[indexes + 1]
keep = np.array(keep)
return keep
推理部分代码：
trt_inference.py
# -*-coding: utf-8-*-
# author: hxy
"""
Inference:
yolov3-tiny.trt
"""
import time
import cv2
from lib import common
import tensorrt as trt
from lib.data_processing import PreprocessYOLO, PostprocessYOLO, VideoPreprocessYOLO
TRT_LOGGER = trt.Logger()
# def load_model(trt_file_path):
# with open(trt_file_path, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime:
# print("Loading engine from: {}".format(trt_file_path.split('/')[-1]))
# return runtime.deserialize_cuda_engine(f.read())
# Inference on Yolov3-tiny
class InferenceYolov3tiny(object):
def __init__(self, engine):
self.engine = engine
# self.img_path = test_img_path
self.input_size = (416, 416)
self.postprocess_args = {"yolo_masks": [(3, 4, 5), (0, 1, 2)],
"yolo_anchors": [(10, 14), (23, 27), (37, 58), (81, 82), (135, 169), (344, 319)],
"obj_threshold": 0.5,
"nms_threshold": 0.35,
"yolo_input_resolution": (416, 416)}
self.output_shape_416 = [(1, 24, 13, 13), (1, 24, 26, 26)]
self.output_shape_480 = [(1, 24, 15, 15), (1, 24, 30, 30)]
self.output_shape_544 = [(1, 24, 17, 17), (1, 24, 34, 34)]
self.output_shape_608 = [(1, 24, 19, 19), (1, 24, 38, 38)]
print("Image input size:{}".format(self.input_size))
def preprocess_img(self, img, context):
# preprocessor = PreprocessYOLO(self.input_size) # 照片处理函数
preprocessor = VideoPreprocessYOLO(self.input_size) # 视频流处理函数
image_raw, image = preprocessor.process(img)
ori_img_hw = image_raw.size
inputs, outputs, bindings, stream = common.allocate_buffers(self.engine)
inputs[0].host = image
s = time.time()
trt_outputs = common.do_inference(context=context,
bindings=bindings,
inputs=inputs,
outputs=outputs,
stream=stream)
ts = time.time() - s
trt_outputs = [output.reshape(shape) for output, shape in zip(trt_outputs, self.output_shape_416)]
postprocessor = PostprocessYOLO(**self.postprocess_args)
boxes, classes, scores = postprocessor.process(trt_outputs, ori_img_hw)
print("Inferences cost: %.3f ms" % ((time.time() - s) * 1000))
if boxes is None:
return image_raw, [], [], [], ts
else:
return image_raw, boxes, classes, scores, ts
这部分代码呢，自己修改的咯，中间ts这个参数，计算的是模型前向推理的时间，便于后面计算FPS
三. 使用tensorrt加速模型推理（视频流）
直接上代码：trt_video_inference.py
# -*-coding: utf-8-*-
# author: hxy
"""
使用tensorrt进行视频流的inference
"""
import cv2
import time
import logging
from PIL import Image
import numpy as np
import tensorrt as trt
from lib.trt_inference import InferenceYolov3tiny
TRT_LOGGER = trt.Logger()
# logging set
def log_set():
logging.basicConfig(level=logging.INFO,
format='%(asctime)s - %(levelname)s - %(message)s')
# 根据预测的id获取类别名
def get_names(names_file, classes_id):
with open(names_file, 'r') as f:
name = f.read()
names = name.splitlines()
f.close()
return names[classes_id]
# 加载model
def load_model(trt_file_path):
with open(trt_file_path, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime:
print("Loading engine: {}!".format(trt_file_path.split('/')[-1]))
return runtime.deserialize_cuda_engine(f.read())
# 获取视频流并进行推理
def video_inference(rtsp, engine, context):
logging.info("获取视频流rtsp地址:{}".format(rtsp))
cap = cv2.VideoCapture('test.mp4')
fourcc = cv2.VideoWriter_fourcc(*'XVID')
size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
out = cv2.VideoWriter('result_trt_inference.avi', fourcc, 10.0, size)
num = 0
while 1:
_, frame = cap.read()
start = time.time()
frame = Image.fromarray(frame, mode="RGB")
image_raw, boxes, classes_id, scores, ts = engine.preprocess_img(img=frame, context=context)
num += 1
print(ts)
fps = 1 / ts
img = cv2.cvtColor(np.asarray(image_raw), cv2.COLOR_RGB2BGR)
for id, box, scores in zip(classes_id, boxes, scores):
name = get_names(names_file='names.txt',
classes_id=id)
logging.info("{}:{}%".format(name, int(scores * 100)))
x_coord, y_coord, width, height = box
left = max(0, np.floor(x_coord + 0.5).astype(int))
top = max(0, np.floor(y_coord + 0.5).astype(int))
right = min(image_raw.width, np.floor(x_coord + width + 0.5).astype(int))
bottom = min(image_raw.height, np.floor(y_coord + height + 0.5).astype(int))
cv2.rectangle(img, (left - 4, top - 4), (right + 4, bottom + 4), (255, 0, 255), 1)
cv2.putText(img, name, (left - 5, top - 8), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
cv2.putText(img, str('FPS: {%.3f}' % fps), (50, 100), cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 255, 0), 2)
img = img[:, :, [2, 1, 0]]
out.write(img)
img = cv2.resize(img, None, fx=.5, fy=.6)
# img = img[:, :, [2, 1, 0]]
cv2.imshow("Inference-Results", img)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
if __name__ == '__main__':
log_set()
with load_model('yolov3-tiny.trt') as engine, engine.create_execution_context() as context:
inference_engine = InferenceYolov3tiny(engine=engine)
# inference_engine = InferenceYolov3(engine=engine)
video_inference(rtsp=" ",
engine=inference_engine,
context=context)
这就是推理代码，里面你需要注意一个地方就是关于names.txt文件，由于本人测试的时候用的模型是只简单训练了一个三类的检测模型，所以呢，我是按照这个来的。这里贴一下本人的names.txt文件内容；
car
pedestrian
face
好了，代码有了 ，那么我们可以开始进行推理了，经过我的软磨硬泡，我朋友的自拍视屏再次被我拉上用场。
python trt_video_inference.py
执行完后，测试视屏的检测结果会输出保存成文件，便于后期查看，至于结果，同样的，我还是截图给大家瞅瞅（朋友再次露脸）
这里呢，这个FPS400多一出来，我都有些不信，但是，代码是自己写的吧，不信也得信呀，不带任何骚操作的推理过程对于一张图片而言确实很快，2-3ms左右。。。对比起原版或者onnx而言，快了好几倍吧。看样子tensorrt的加速效果还是可以的吧。
总结
- 本篇博客采用tensorrt来加速yolov3-tiny网络框架的推理，从用来测试的mp4文件上的fps以及单张图片的前向推理时间来看，很快，很猛！
- 当然，并不是所有的代码都是自己写的，参考加学习！！
- 所以，文中有不足支出，还请多多谅解！ 要是有内容描述不对的还请指出，以便共同学习，进步！！
- 希望本博客对您有用！！