├── README.txt ├── common.py ├── data_processing.py ├── images ├── 1.PNG ├── 1.jpg ├── 10.PNG ├── 10.jpg ├── 11.PNG ├── 11.jpg ├── 12.PNG ├── 12.jpg ├── 13.PNG ├── 13.jpg ├── 2.PNG ├── 2.jpg ├── 3.PNG ├── 3.jpg ├── 4.PNG ├── 4.jpg ├── 5.PNG ├── 5.jpg ├── 6.PNG ├── 6.jpg ├── 7.PNG ├── 7.jpg ├── 8.PNG ├── 8.jpg ├── 9.PNG ├── 9.jpg └── PNG2jpg.py ├── inference.py └── requirements.txt /README.txt: -------------------------------------------------------------------------------- 1 | 一、简介 2 | 本文档主要介绍如何在Nvidia tx2上使用tensorrt加速yolov3-tiny。 3 | 二、步骤 4 | 1、环境配置 5 | 1.1 所需环境： jetpack4.2: 6 | ubuntu 18.04 7 | tensorrt 5.0.6.3 8 | cuda 10.0 9 | cudnn 7.3.1 10 | 1.2 tx2安装环境较麻烦，建议使用jetpack刷机，将自动安装最新驱动、CUDA Toolkit、cuDNN、tensorRT等所需软件，一步到位。 11 | 下面以jetpack4.2为例，说明tx2刷机步骤。 12 | (1)所需设备：NVIDIA tx2及其配件 x1; 13 | 显示器(连接tx2) x1; 14 | 键盘(连接tx2) x1; 15 | 鼠标(连接tx2) x1; 16 | host机器 x1; 17 | 网线 x2; 18 | 说明：1)host主机需要x86_64结构的ubuntu16操作系统，用来运行sdkmanager-[version].[build#].deb文件（实测可以使用虚拟机）; 19 | 2)烧写期间需要将tx2与host连接到同一局域网内; 20 | 3)烧写jetpack4.2需要在tx2端操作，需提前准备鼠标、键盘以及显示器; 21 | 4)烧写时会从互联网下载各种包，网络原因可能导致无法连接至安装源; 22 | 5)烧写耗时较长; 23 | (2)准备 sdkmanager-[version].[build#].deb (见压缩包根目录) 24 | (3)在host机器上运行.deb文件，进入jetpack4.2可视界面; 25 | (4)根据提示，选择设备型号等; 26 | (5)进入下载页面，待install进程到40%时，会提示连接tx2，选择manual step; 27 | (6)之后将设备置于usb恢复模式，所以 28 | 1)断开tx2电源，使其处于断电关机状态; 29 | 2)用网线将tx2连入host所在的局域网中; 30 | 3)用Micro USB线(黑色)将tx2连接到host机上; 31 | 4)接通tx2电源，按power键开机; 32 | 5)长按recovery键不松开，同时点按一下reset键，等待3s后松开recovery键，此时tx2处于强制恢复模式。 33 | (7)检查tx2与host是否正确连接，打开终端输入lsusb，列表中有Nvidia Corp说明连接成功; 34 | (8)检查无误后点击flash; 35 | (9)当Install到50%时tx2会开机，根据指示设置完成ubuntu18系统，在host端输入账号密码，开始安装软件; 36 | (10)当install到100%时，刷机完成。 37 | 38 | 刷机参考： 39 | https://blog.csdn.net/weixin_43842032/article/details/88753724 40 | 2、使用说明 41 | 2.1 依赖安装及脚本使用 42 | (1)安装所需软件包 43 | 解压 yolov3-tiny_trt.rar 44 | 进入yolov3-tiny_trt文件夹 45 | 使用安装requirement.txt中所列的包： 46 | # 使用pip安装，需先安装pip： 47 | sudo apt-get install python-pip 48 | 49 | # 安装pycuda等： 50 | pip install pycuda 51 | pip install Pillow # 安装Pillow时可能会提示缺少依赖 52 | 53 | (2)脚本使用： 54 | # 打开inference.py，在105行至107行处分别指定 输入尺寸，输入图像列表文件，保存路径： 55 | INPUT_SIZE = 608 # [416, 480, 544, 608] 56 | INPUT_LIST_FILE = './ped_list.txt' 57 | SAVE_PATH = './images_results/' 58 | 59 | # 之后运行脚本 60 | python inference.py 61 | 62 | 可得到检测结果，检测时间，并保存图像至指定目录。 63 | 64 | Performance：engine_416: 9.486ms 65 | engine_480: 11.174ms 66 | engine_544: 17.271ms 67 | engine_608: 20.249ms 68 | 2.2 使用说明 69 | 在tx2上运行脚本时，建议开始最大功耗模式，并开启风扇，以实现最优加速效果： 70 | # 查询当前工作模式 71 | sudo nvpmodel -q verbose 72 | 73 | # 开启最大工作模式 MAXN 0 74 | sudo nvpmodel -m 0 75 | 76 | # 打开风扇 77 | sudo /usr/bin/jetson_clocks 78 | 79 | 3、工程结构与调用说明 80 | 3.1 文件列表： 81 | --yolov3-tiny_trt/ 82 | --engines/ 83 | --images/ 84 | --images_results/ 85 | --README.txt 86 | --requirements.txt 87 | --common.py 88 | --data_processing.py 89 | --inference.py 90 | --ped_list.txt 91 | 3.2 (1)如3.1，engines中保存有416, 480, 544, 608, 4种不同输入尺寸的engine文件，用于推理; 92 | (2)在images中存数张有测试图像; 93 | (3)requirements.txt中为所需包的列表; 94 | (4)common.py文件中，定义allocate_buffers()函数以及do_inference()函数，用于分配空间及推理; 95 | (5)data_processing.py中，分别定义函数，用于图像的加载和预处理，以及模型输出的后处理，如nms等; 96 | (6)inference.py中，定义检测框绘制函数和engine加载函数，以及检测main函数，加载指定engine文件后，对输入进行推理，后对输出结果进行处理，绘制检测结果并保存。 97 | 98 | 99 | 100 | -------------------------------------------------------------------------------- /common.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright 1993-2018 NVIDIA Corporation. All rights reserved. 3 | # 4 | # NOTICE TO LICENSEE: 5 | # 6 | # This source code and/or documentation ("Licensed Deliverables") are 7 | # subject to NVIDIA intellectual property rights under U.S. and 8 | # international Copyright laws. 9 | # 10 | # These Licensed Deliverables contained herein is PROPRIETARY and 11 | # CONFIDENTIAL to NVIDIA and is being provided under the terms and 12 | # conditions of a form of NVIDIA software license agreement by and 13 | # between NVIDIA and Licensee ("License Agreement") or electronically 14 | # accepted by Licensee. Notwithstanding any terms or conditions to 15 | # the contrary in the License Agreement, reproduction or disclosure 16 | # of the Licensed Deliverables to any third party without the express 17 | # written consent of NVIDIA is prohibited. 18 | # 19 | # NOTWITHSTANDING ANY TERMS OR CONDITIONS TO THE CONTRARY IN THE 20 | # LICENSE AGREEMENT, NVIDIA MAKES NO REPRESENTATION ABOUT THE 21 | # SUITABILITY OF THESE LICENSED DELIVERABLES FOR ANY PURPOSE. IT IS 22 | # PROVIDED "AS IS" WITHOUT EXPRESS OR IMPLIED WARRANTY OF ANY KIND. 23 | # NVIDIA DISCLAIMS ALL WARRANTIES WITH REGARD TO THESE LICENSED 24 | # DELIVERABLES, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY, 25 | # NONINFRINGEMENT, AND FITNESS FOR A PARTICULAR PURPOSE. 26 | # NOTWITHSTANDING ANY TERMS OR CONDITIONS TO THE CONTRARY IN THE 27 | # LICENSE AGREEMENT, IN NO EVENT SHALL NVIDIA BE LIABLE FOR ANY 28 | # SPECIAL, INDIRECT, INCIDENTAL, OR CONSEQUENTIAL DAMAGES, OR ANY 29 | # DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, 30 | # WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS 31 | # ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE 32 | # OF THESE LICENSED DELIVERABLES. 33 | # 34 | # U.S. Government End Users. These Licensed Deliverables are a 35 | # "commercial item" as that term is defined at 48 C.F.R. 2.101 (OCT 36 | # 1995), consisting of "commercial computer software" and "commercial 37 | # computer software documentation" as such terms are used in 48 38 | # C.F.R. 12.212 (SEPT 1995) and is provided to the U.S. Government 39 | # only as a commercial end item. Consistent with 48 C.F.R.12.212 and 40 | # 48 C.F.R. 227.7202-1 through 227.7202-4 (JUNE 1995), all 41 | # U.S. Government End Users acquire the Licensed Deliverables with 42 | # only those rights set forth herein. 43 | # 44 | # Any use of the Licensed Deliverables in individual and commercial 45 | # software must include, in the user documentation and internal 46 | # comments to the code, the above Disclaimer and U.S. Government End 47 | # Users Notice. 48 | # 49 | 50 | import os 51 | import argparse 52 | import numpy as np 53 | import pycuda.driver as cuda 54 | import tensorrt as trt 55 | 56 | try: 57 | # Sometimes python2 does not understand FileNotFoundError 58 | FileNotFoundError 59 | except NameError: 60 | FileNotFoundError = IOError 61 | 62 | def GiB(val): 63 | return val * 1 << 30 64 | 65 | def find_sample_data(description="Runs a TensorRT Python sample", subfolder="", find_files=[]): 66 | ''' 67 | Parses sample arguments. 68 | Args: 69 | description (str): Description of the sample. 70 | subfolder (str): The subfolder containing data relevant to this sample 71 | find_files (str): A list of filenames to find. Each filename will be replaced with an absolute path. 72 | Returns: 73 | str: Path of data directory. 74 | Raises: 75 | FileNotFoundError 76 | ''' 77 | kDEFAULT_DATA_ROOT = os.path.abspath("/usr/src/tensorrt/data") 78 | 79 | # Standard command-line arguments for all samples. 80 | parser = argparse.ArgumentParser(description=description) 81 | parser.add_argument("-d", "--datadir", help="Location of the TensorRT sample data directory.") 82 | args, unknown_args = parser.parse_known_args() 83 | 84 | # If data directory is not specified, use the default. 85 | data_root = args.datadir if args.datadir else kDEFAULT_DATA_ROOT 86 | # If the subfolder exists, append it to the path, otherwise use the provided path as-is. 87 | subfolder_path = os.path.join(data_root, subfolder) 88 | if not os.path.exists(subfolder_path): 89 | print("WARNING: " + subfolder_path + " does not exist. Using " + data_root + " instead.") 90 | data_path = subfolder_path if os.path.exists(subfolder_path) else data_root 91 | 92 | # Make sure data directory exists. 93 | if not (os.path.exists(data_path)): 94 | raise FileNotFoundError(data_path + " does not exist. Please provide the correct data path with the -d option.") 95 | 96 | # Find all requested files. 97 | for index, f in enumerate(find_files): 98 | find_files[index] = os.path.abspath(os.path.join(data_path, f)) 99 | if not os.path.exists(find_files[index]): 100 | raise FileNotFoundError(find_files[index] + " does not exist. Please provide the correct data path with the -d option.") 101 | if find_files: 102 | return data_path, find_files 103 | else: 104 | return data_path 105 | 106 | # Simple helper data class that's a little nicer to use than a 2-tuple. 107 | class HostDeviceMem(object): 108 | def __init__(self, host_mem, device_mem): 109 | self.host = host_mem 110 | self.device = device_mem 111 | 112 | def __str__(self): 113 | return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device) 114 | 115 | def __repr__(self): 116 | return self.__str__() 117 | 118 | # Allocates all buffers required for an engine, i.e. host/device inputs/outputs. 119 | def allocate_buffers(engine): 120 | inputs = [] 121 | outputs = [] 122 | bindings = [] 123 | stream = cuda.Stream() 124 | for binding in engine: 125 | size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size 126 | dtype = trt.nptype(engine.get_binding_dtype(binding)) 127 | # Allocate host and device buffers 128 | host_mem = cuda.pagelocked_empty(size, dtype) 129 | device_mem = cuda.mem_alloc(host_mem.nbytes) 130 | # Append the device buffer to device bindings. 131 | bindings.append(int(device_mem)) 132 | # Append to the appropriate list. 133 | if engine.binding_is_input(binding): 134 | inputs.append(HostDeviceMem(host_mem, device_mem)) 135 | else: 136 | outputs.append(HostDeviceMem(host_mem, device_mem)) 137 | return inputs, outputs, bindings, stream 138 | 139 | # This function is generalized for multiple inputs/outputs. 140 | # inputs and outputs are expected to be lists of HostDeviceMem objects. 141 | def do_inference(context, bindings, inputs, outputs, stream, batch_size=1): 142 | # Transfer input data to the GPU. 143 | [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs] 144 | # Run inference. 145 | context.execute_async(batch_size=batch_size, bindings=bindings, stream_handle=stream.handle) 146 | # Transfer predictions back from the GPU. 147 | [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs] 148 | # Synchronize the stream 149 | stream.synchronize() 150 | # Return only the host outputs. 151 | return [out.host for out in outputs] 152 | -------------------------------------------------------------------------------- /data_processing.py: -------------------------------------------------------------------------------- 1 | import math 2 | from PIL import Image 3 | import numpy as np 4 | 5 | ALL_CATEGORIES = ['person'] 6 | CATEGORY_NUM = len(ALL_CATEGORIES) 7 | assert CATEGORY_NUM == 1 8 | 9 | output_shapes_416 = [(1, 18, 13, 13), (1, 18, 26, 26)] 10 | output_shapes_480 = [(1, 18, 15, 15), (1, 18, 30, 30)] 11 | output_shapes_544 = [(1, 18, 17, 17), (1, 18, 34, 34)] 12 | output_shapes_608 = [(1, 18, 19, 19), (1, 18, 38, 38)] 13 | OUTPUT_SHAPES_DICT = {'416': output_shapes_416, '480': output_shapes_480, '544': output_shapes_544, '608': output_shapes_608} 14 | 15 | 16 | class PreprocessYOLO(object): 17 | """A simple class for loading images with PIL and reshaping them to the specified 18 | input resolution for YOLOv3-608. 19 | """ 20 | 21 | def __init__(self, yolo_input_resolution): 22 | """Initialize with the input resolution for YOLOv3, which will stay fixed in this sample. 23 | 24 | Keyword arguments: 25 | yolo_input_resolution -- two-dimensional tuple with the target network's (spatial) 26 | input resolution in HW order 27 | """ 28 | self.yolo_input_resolution = yolo_input_resolution 29 | 30 | def process(self, input_image_path): 31 | """Load an image from the specified input path, 32 | and return it together with a pre-processed version required for feeding it into a 33 | YOLOv3 network. 34 | 35 | Keyword arguments: 36 | input_image_path -- string path of the image to be loaded 37 | """ 38 | image_raw, image_resized = self._load_and_resize(input_image_path) 39 | image_preprocessed = self._shuffle_and_normalize(image_resized) 40 | return image_raw, image_preprocessed 41 | 42 | def _load_and_resize(self, input_image_path): 43 | """Load an image from the specified path and resize it to the input resolution. 44 | Return the input image before resizing as a PIL Image (required for visualization), 45 | and the resized image as a NumPy float array. 46 | 47 | Keyword arguments: 48 | input_image_path -- string path of the image to be loaded 49 | """ 50 | 51 | image_raw = Image.open(input_image_path) 52 | # Expecting yolo_input_resolution in (height, width) format, adjusting to PIL 53 | # convention (width, height) in PIL: 54 | new_resolution = ( 55 | self.yolo_input_resolution[1], 56 | self.yolo_input_resolution[0]) 57 | print(new_resolution) 58 | image_resized = image_raw.resize( 59 | new_resolution, resample=Image.BICUBIC) 60 | image_resized = np.array(image_resized, dtype=np.float32, order='C') 61 | return image_raw, image_resized 62 | 63 | def _shuffle_and_normalize(self, image): 64 | """Normalize a NumPy array representing an image to the range [0, 1], and 65 | convert it from HWC format ("channels last") to NCHW format ("channels first" 66 | with leading batch dimension). 67 | 68 | Keyword arguments: 69 | image -- image as three-dimensional NumPy float array, in HWC format 70 | """ 71 | image /= 255.0 72 | # HWC to CHW format: 73 | image = np.transpose(image, [2, 0, 1]) 74 | # CHW to NCHW format 75 | image = np.expand_dims(image, axis=0) 76 | # Convert the image to row-major order, also known as "C order": 77 | image = np.array(image, dtype=np.float32, order='C') 78 | return image 79 | 80 | 81 | class PostprocessYOLO(object): 82 | """Class for post-processing the three outputs tensors from YOLOv3-608.""" 83 | 84 | def __init__(self, 85 | yolo_masks, 86 | yolo_anchors, 87 | obj_threshold, 88 | nms_threshold, 89 | yolo_input_resolution): 90 | """Initialize with all values that will be kept when processing several frames. 91 | Assuming 3 outputs of the network in the case of (large) YOLOv3. 92 | 93 | Keyword arguments: 94 | yolo_masks -- a list of 3 three-dimensional tuples for the YOLO masks 95 | yolo_anchors -- a list of 9 two-dimensional tuples for the YOLO anchors 96 | object_threshold -- threshold for object coverage, float value between 0 and 1 97 | nms_threshold -- threshold for non-max suppression algorithm, 98 | float value between 0 and 1 99 | input_resolution_yolo -- two-dimensional tuple with the target network's (spatial) 100 | input resolution in HW order 101 | """ 102 | self.masks = yolo_masks 103 | self.anchors = yolo_anchors 104 | self.object_threshold = obj_threshold 105 | self.nms_threshold = nms_threshold 106 | self.input_resolution_yolo = yolo_input_resolution 107 | 108 | def process(self, outputs, resolution_raw): 109 | """Take the YOLOv3 outputs generated from a TensorRT forward pass, post-process them 110 | and return a list of bounding boxes for detected object together with their category 111 | and their confidences in separate lists. 112 | 113 | Keyword arguments: 114 | outputs -- outputs from a TensorRT engine in NCHW format 115 | resolution_raw -- the original spatial resolution from the input PIL image in WH order 116 | """ 117 | outputs_reshaped = list() 118 | for output in outputs: 119 | outputs_reshaped.append(self._reshape_output(output)) 120 | 121 | boxes, categories, confidences = self._process_yolo_output( 122 | outputs_reshaped, resolution_raw) 123 | 124 | return boxes, categories, confidences 125 | 126 | def _reshape_output(self, output): 127 | """Reshape a TensorRT output from NCHW to NHWC format (with expected C=255), 128 | and then return it in (height,width,3,85) dimensionality after further reshaping. 129 | 130 | Keyword argument: 131 | output -- an output from a TensorRT engine after inference 132 | """ 133 | output = np.transpose(output, [0, 2, 3, 1]) 134 | _, height, width, _ = output.shape 135 | dim1, dim2 = height, width 136 | dim3 = 3 137 | # There are CATEGORY_NUM=80 object categories: 138 | dim4 = (4 + 1 + CATEGORY_NUM) 139 | return np.reshape(output, (dim1, dim2, dim3, dim4)) 140 | 141 | def _process_yolo_output(self, outputs_reshaped, resolution_raw): 142 | """Take in a list of three reshaped YOLO outputs in (height,width,3,85) shape and return 143 | return a list of bounding boxes for detected object together with their category and their 144 | confidences in separate lists. 145 | 146 | Keyword arguments: 147 | outputs_reshaped -- list of three reshaped YOLO outputs as NumPy arrays 148 | with shape (height,width,3,85) 149 | resolution_raw -- the original spatial resolution from the input PIL image in WH order 150 | """ 151 | 152 | # E.g. in YOLOv3-608, there are three output tensors, which we associate with their 153 | # respective masks. Then we iterate through all output-mask pairs and generate candidates 154 | # for bounding boxes, their corresponding category predictions and their confidences: 155 | boxes, categories, confidences = list(), list(), list() 156 | for output, mask in zip(outputs_reshaped, self.masks): 157 | box, category, confidence = self._process_feats(output, mask) 158 | box, category, confidence = self._filter_boxes(box, category, confidence) 159 | boxes.append(box) 160 | categories.append(category) 161 | confidences.append(confidence) 162 | 163 | boxes = np.concatenate(boxes) 164 | categories = np.concatenate(categories) 165 | confidences = np.concatenate(confidences) 166 | 167 | # Scale boxes back to original image shape: 168 | width, height = resolution_raw 169 | image_dims = [width, height, width, height] 170 | boxes = boxes * image_dims 171 | 172 | # Using the candidates from the previous (loop) step, we apply the non-max suppression 173 | # algorithm that clusters adjacent bounding boxes to a single bounding box: 174 | nms_boxes, nms_categories, nscores = list(), list(), list() 175 | for category in set(categories): 176 | idxs = np.where(categories == category) 177 | box = boxes[idxs] 178 | category = categories[idxs] 179 | confidence = confidences[idxs] 180 | 181 | keep = self._nms_boxes(box, confidence) 182 | 183 | nms_boxes.append(box[keep]) 184 | nms_categories.append(category[keep]) 185 | nscores.append(confidence[keep]) 186 | 187 | if not nms_categories and not nscores: 188 | return [], [], [] 189 | 190 | boxes = np.concatenate(nms_boxes) 191 | categories = np.concatenate(nms_categories) 192 | confidences = np.concatenate(nscores) 193 | 194 | return boxes, categories, confidences 195 | 196 | def _process_feats(self, output_reshaped, mask): 197 | """Take in a reshaped YOLO output in height,width,3,85 format together with its 198 | corresponding YOLO mask and return the detected bounding boxes, the confidence, 199 | and the class probability in each cell/pixel. 200 | 201 | Keyword arguments: 202 | output_reshaped -- reshaped YOLO output as NumPy arrays with shape (height,width,3,85) 203 | mask -- 2-dimensional tuple with mask specification for this output 204 | """ 205 | 206 | # Two in-line functions required for calculating the bounding box 207 | # descriptors: 208 | def sigmoid(value): 209 | """Return the sigmoid of the input.""" 210 | return 1.0 / (1.0 + math.exp(-value)) 211 | 212 | def exponential(value): 213 | """Return the exponential of the input.""" 214 | return math.exp(value) 215 | 216 | # Vectorized calculation of above two functions: 217 | sigmoid_v = np.vectorize(sigmoid) 218 | exponential_v = np.vectorize(exponential) 219 | 220 | grid_h, grid_w, _, _ = output_reshaped.shape 221 | 222 | anchors = [self.anchors[i] for i in mask] 223 | 224 | # Reshape to N, height, width, num_anchors, box_params: 225 | anchors_tensor = np.reshape(anchors, [1, 1, len(anchors), 2]) 226 | box_xy = sigmoid_v(output_reshaped[..., :2]) 227 | box_wh = exponential_v(output_reshaped[..., 2:4]) * anchors_tensor 228 | box_confidence = sigmoid_v(output_reshaped[..., 4]) 229 | 230 | box_confidence = np.expand_dims(box_confidence, axis=-1) 231 | box_class_probs = sigmoid_v(output_reshaped[..., 5:]) 232 | 233 | col = np.tile(np.arange(0, grid_w), grid_w).reshape(-1, grid_w) 234 | row = np.tile(np.arange(0, grid_h).reshape(-1, 1), grid_h) 235 | 236 | col = col.reshape(grid_h, grid_w, 1, 1).repeat(3, axis=-2) 237 | row = row.reshape(grid_h, grid_w, 1, 1).repeat(3, axis=-2) 238 | grid = np.concatenate((col, row), axis=-1) 239 | 240 | box_xy += grid 241 | box_xy /= (grid_w, grid_h) 242 | box_wh /= self.input_resolution_yolo 243 | box_xy -= (box_wh / 2.) 244 | boxes = np.concatenate((box_xy, box_wh), axis=-1) 245 | 246 | # boxes: centroids, box_confidence: confidence level, box_class_probs: 247 | # class confidence 248 | return boxes, box_confidence, box_class_probs 249 | 250 | def _filter_boxes(self, boxes, box_confidences, box_class_probs): 251 | """Take in the unfiltered bounding box descriptors and discard each cell 252 | whose score is lower than the object threshold set during class initialization. 253 | 254 | Keyword arguments: 255 | boxes -- bounding box coordinates with shape (height,width,3,4); 4 for 256 | x,y,height,width coordinates of the boxes 257 | box_confidences -- bounding box confidences with shape (height,width,3,1); 1 for as 258 | confidence scalar per element 259 | box_class_probs -- class probabilities with shape (height,width,3,CATEGORY_NUM) 260 | 261 | """ 262 | box_scores = box_confidences * box_class_probs 263 | box_classes = np.argmax(box_scores, axis=-1) 264 | box_class_scores = np.max(box_scores, axis=-1) 265 | pos = np.where(box_class_scores >= self.object_threshold) 266 | 267 | boxes = boxes[pos] 268 | classes = box_classes[pos] 269 | scores = box_class_scores[pos] 270 | 271 | return boxes, classes, scores 272 | 273 | def _nms_boxes(self, boxes, box_confidences): 274 | """Apply the Non-Maximum Suppression (NMS) algorithm on the bounding boxes with their 275 | confidence scores and return an array with the indexes of the bounding boxes we want to 276 | keep (and display later). 277 | 278 | Keyword arguments: 279 | boxes -- a NumPy array containing N bounding-box coordinates that survived filtering, 280 | with shape (N,4); 4 for x,y,height,width coordinates of the boxes 281 | box_confidences -- a Numpy array containing the corresponding confidences with shape N 282 | """ 283 | x_coord = boxes[:, 0] 284 | y_coord = boxes[:, 1] 285 | width = boxes[:, 2] 286 | height = boxes[:, 3] 287 | 288 | areas = width * height 289 | ordered = box_confidences.argsort()[::-1] 290 | 291 | keep = list() 292 | while ordered.size > 0: 293 | # Index of the current element: 294 | i = ordered[0] 295 | keep.append(i) 296 | xx1 = np.maximum(x_coord[i], x_coord[ordered[1:]]) 297 | yy1 = np.maximum(y_coord[i], y_coord[ordered[1:]]) 298 | xx2 = np.minimum(x_coord[i] + width[i], x_coord[ordered[1:]] + width[ordered[1:]]) 299 | yy2 = np.minimum(y_coord[i] + height[i], y_coord[ordered[1:]] + height[ordered[1:]]) 300 | 301 | width1 = np.maximum(0.0, xx2 - xx1 + 1) 302 | height1 = np.maximum(0.0, yy2 - yy1 + 1) 303 | intersection = width1 * height1 304 | union = (areas[i] + areas[ordered[1:]] - intersection) 305 | 306 | # Compute the Intersection over Union (IoU) score: 307 | iou = intersection / union 308 | 309 | # The goal of the NMS algorithm is to reduce the number of adjacent bounding-box 310 | # candidates to a minimum. In this step, we keep only those elements whose overlap 311 | # with the current bounding box is lower than the threshold: 312 | indexes = np.where(iou <= self.nms_threshold)[0] 313 | ordered = ordered[indexes + 1] 314 | 315 | keep = np.array(keep) 316 | return keep 317 | -------------------------------------------------------------------------------- /images/1.PNG: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CaoE/tensorrt_yolov3-tiny/18beeb7faf37dc1edfdc7a87a70b0897fc2fc0a9/images/1.PNG -------------------------------------------------------------------------------- /images/1.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CaoE/tensorrt_yolov3-tiny/18beeb7faf37dc1edfdc7a87a70b0897fc2fc0a9/images/1.jpg -------------------------------------------------------------------------------- /images/10.PNG: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CaoE/tensorrt_yolov3-tiny/18beeb7faf37dc1edfdc7a87a70b0897fc2fc0a9/images/10.PNG -------------------------------------------------------------------------------- /images/10.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CaoE/tensorrt_yolov3-tiny/18beeb7faf37dc1edfdc7a87a70b0897fc2fc0a9/images/10.jpg -------------------------------------------------------------------------------- /images/11.PNG: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CaoE/tensorrt_yolov3-tiny/18beeb7faf37dc1edfdc7a87a70b0897fc2fc0a9/images/11.PNG -------------------------------------------------------------------------------- /images/11.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CaoE/tensorrt_yolov3-tiny/18beeb7faf37dc1edfdc7a87a70b0897fc2fc0a9/images/11.jpg -------------------------------------------------------------------------------- /images/12.PNG: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CaoE/tensorrt_yolov3-tiny/18beeb7faf37dc1edfdc7a87a70b0897fc2fc0a9/images/12.PNG -------------------------------------------------------------------------------- /images/12.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CaoE/tensorrt_yolov3-tiny/18beeb7faf37dc1edfdc7a87a70b0897fc2fc0a9/images/12.jpg -------------------------------------------------------------------------------- /images/13.PNG: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CaoE/tensorrt_yolov3-tiny/18beeb7faf37dc1edfdc7a87a70b0897fc2fc0a9/images/13.PNG -------------------------------------------------------------------------------- /images/13.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CaoE/tensorrt_yolov3-tiny/18beeb7faf37dc1edfdc7a87a70b0897fc2fc0a9/images/13.jpg -------------------------------------------------------------------------------- /images/2.PNG: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CaoE/tensorrt_yolov3-tiny/18beeb7faf37dc1edfdc7a87a70b0897fc2fc0a9/images/2.PNG -------------------------------------------------------------------------------- /images/2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CaoE/tensorrt_yolov3-tiny/18beeb7faf37dc1edfdc7a87a70b0897fc2fc0a9/images/2.jpg -------------------------------------------------------------------------------- /images/3.PNG: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CaoE/tensorrt_yolov3-tiny/18beeb7faf37dc1edfdc7a87a70b0897fc2fc0a9/images/3.PNG -------------------------------------------------------------------------------- /images/3.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CaoE/tensorrt_yolov3-tiny/18beeb7faf37dc1edfdc7a87a70b0897fc2fc0a9/images/3.jpg -------------------------------------------------------------------------------- /images/4.PNG: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CaoE/tensorrt_yolov3-tiny/18beeb7faf37dc1edfdc7a87a70b0897fc2fc0a9/images/4.PNG -------------------------------------------------------------------------------- /images/4.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CaoE/tensorrt_yolov3-tiny/18beeb7faf37dc1edfdc7a87a70b0897fc2fc0a9/images/4.jpg -------------------------------------------------------------------------------- /images/5.PNG: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CaoE/tensorrt_yolov3-tiny/18beeb7faf37dc1edfdc7a87a70b0897fc2fc0a9/images/5.PNG -------------------------------------------------------------------------------- /images/5.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CaoE/tensorrt_yolov3-tiny/18beeb7faf37dc1edfdc7a87a70b0897fc2fc0a9/images/5.jpg -------------------------------------------------------------------------------- /images/6.PNG: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CaoE/tensorrt_yolov3-tiny/18beeb7faf37dc1edfdc7a87a70b0897fc2fc0a9/images/6.PNG -------------------------------------------------------------------------------- /images/6.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CaoE/tensorrt_yolov3-tiny/18beeb7faf37dc1edfdc7a87a70b0897fc2fc0a9/images/6.jpg -------------------------------------------------------------------------------- /images/7.PNG: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CaoE/tensorrt_yolov3-tiny/18beeb7faf37dc1edfdc7a87a70b0897fc2fc0a9/images/7.PNG -------------------------------------------------------------------------------- /images/7.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CaoE/tensorrt_yolov3-tiny/18beeb7faf37dc1edfdc7a87a70b0897fc2fc0a9/images/7.jpg -------------------------------------------------------------------------------- /images/8.PNG: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CaoE/tensorrt_yolov3-tiny/18beeb7faf37dc1edfdc7a87a70b0897fc2fc0a9/images/8.PNG -------------------------------------------------------------------------------- /images/8.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CaoE/tensorrt_yolov3-tiny/18beeb7faf37dc1edfdc7a87a70b0897fc2fc0a9/images/8.jpg -------------------------------------------------------------------------------- /images/9.PNG: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CaoE/tensorrt_yolov3-tiny/18beeb7faf37dc1edfdc7a87a70b0897fc2fc0a9/images/9.PNG -------------------------------------------------------------------------------- /images/9.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CaoE/tensorrt_yolov3-tiny/18beeb7faf37dc1edfdc7a87a70b0897fc2fc0a9/images/9.jpg -------------------------------------------------------------------------------- /images/PNG2jpg.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | import numpy as np 3 | import cv2 4 | from PIL import ImageDraw 5 | from PIL import Image 6 | import os 7 | import sys 8 | 9 | path='./' 10 | 11 | for filename in os.listdir(path): 12 | if os.path.splitext(filename)[1]=='.PNG': 13 | img = cv2.imread(path + filename) 14 | new_filename = filename.replace('.PNG', '.jpg') 15 | cv2.imwrite(path + new_filename, img) 16 | 17 | 18 | #image = cv2.imread('./2.PNG') 19 | #cv2.imwrite('./2.jpg', image) 20 | -------------------------------------------------------------------------------- /inference.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python2 2 | # 3 | # Copyright 1993-2018 NVIDIA Corporation. All rights reserved. 4 | # 5 | # NOTICE TO LICENSEE: 6 | # 7 | # This source code and/or documentation ("Licensed Deliverables") are 8 | # subject to NVIDIA intellectual property rights under U.S. and 9 | # international Copyright laws. 10 | # 11 | # These Licensed Deliverables contained herein is PROPRIETARY and 12 | # CONFIDENTIAL to NVIDIA and is being provided under the terms and 13 | # conditions of a form of NVIDIA software license agreement by and 14 | # between NVIDIA and Licensee ("License Agreement") or electronically 15 | # accepted by Licensee. Notwithstanding any terms or conditions to 16 | # the contrary in the License Agreement, reproduction or disclosure 17 | # of the Licensed Deliverables to any third party without the express 18 | # written consent of NVIDIA is prohibited. 19 | # 20 | # NOTWITHSTANDING ANY TERMS OR CONDITIONS TO THE CONTRARY IN THE 21 | # LICENSE AGREEMENT, NVIDIA MAKES NO REPRESENTATION ABOUT THE 22 | # SUITABILITY OF THESE LICENSED DELIVERABLES FOR ANY PURPOSE. IT IS 23 | # PROVIDED "AS IS" WITHOUT EXPRESS OR IMPLIED WARRANTY OF ANY KIND. 24 | # NVIDIA DISCLAIMS ALL WARRANTIES WITH REGARD TO THESE LICENSED 25 | # DELIVERABLES, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY, 26 | # NONINFRINGEMENT, AND FITNESS FOR A PARTICULAR PURPOSE. 27 | # NOTWITHSTANDING ANY TERMS OR CONDITIONS TO THE CONTRARY IN THE 28 | # LICENSE AGREEMENT, IN NO EVENT SHALL NVIDIA BE LIABLE FOR ANY 29 | # SPECIAL, INDIRECT, INCIDENTAL, OR CONSEQUENTIAL DAMAGES, OR ANY 30 | # DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, 31 | # WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS 32 | # ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE 33 | # OF THESE LICENSED DELIVERABLES. 34 | # 35 | # U.S. Government End Users. These Licensed Deliverables are a 36 | # "commercial item" as that term is defined at 48 C.F.R. 2.101 (OCT 37 | # 1995), consisting of "commercial computer software" and "commercial 38 | # computer software documentation" as such terms are used in 48 39 | # C.F.R. 12.212 (SEPT 1995) and is provided to the U.S. Government 40 | # only as a commercial end item. Consistent with 48 C.F.R.12.212 and 41 | # 48 C.F.R. 227.7202-1 through 227.7202-4 (JUNE 1995), all 42 | # U.S. Government End Users acquire the Licensed Deliverables with 43 | # only those rights set forth herein. 44 | # 45 | # Any use of the Licensed Deliverables in individual and commercial 46 | # software must include, in the user documentation and internal 47 | # comments to the code, the above Disclaimer and U.S. Government End 48 | # Users Notice. 49 | # 50 | 51 | from __future__ import print_function 52 | 53 | import time 54 | import numpy as np 55 | import tensorrt as trt 56 | import pycuda.driver as cuda 57 | import pycuda.autoinit 58 | from PIL import ImageDraw 59 | 60 | from data_processing import PreprocessYOLO, PostprocessYOLO, ALL_CATEGORIES, \ 61 | OUTPUT_SHAPES_DICT 62 | 63 | import sys, os 64 | import common 65 | 66 | 67 | TRT_LOGGER = trt.Logger() 68 | 69 | def draw_bboxes(image_raw, bboxes, confidences, categories, all_categories, bbox_color='red'): 70 | """Draw the bounding boxes on the original input image and return it. 71 | 72 | Keyword arguments: 73 | image_raw -- a raw PIL Image 74 | bboxes -- NumPy array containing the bounding box coordinates of N objects, with shape (N,4). 75 | categories -- NumPy array containing the corresponding category for each object, 76 | with shape (N,) 77 | confidences -- NumPy array containing the corresponding confidence for each object, 78 | with shape (N,) 79 | all_categories -- a list of all categories in the correct ordered (required for looking up 80 | the category name) 81 | bbox_color -- an optional string specifying the color of the bounding boxes (default: 'blue') 82 | """ 83 | draw = ImageDraw.Draw(image_raw) 84 | print(bboxes, confidences, categories) 85 | for box, score, category in zip(bboxes, confidences, categories): 86 | x_coord, y_coord, width, height = box 87 | left = max(0, np.floor(x_coord + 0.5).astype(int)) 88 | top = max(0, np.floor(y_coord + 0.5).astype(int)) 89 | right = min(image_raw.width, np.floor(x_coord + width + 0.5).astype(int)) 90 | bottom = min(image_raw.height, np.floor(y_coord + height + 0.5).astype(int)) 91 | 92 | draw.rectangle(((left, top), (right, bottom)), outline=bbox_color) 93 | draw.text((left, top - 12), '{0} {1:.2f}'.format(all_categories[category], score), fill=bbox_color) 94 | 95 | return image_raw 96 | 97 | def get_engine(engine_file_path=""): 98 | print("Reading engine from file {}".format(engine_file_path)) 99 | with open(engine_file_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime: 100 | return runtime.deserialize_cuda_engine(f.read()) 101 | 102 | def main(): 103 | 104 | # configure 105 | INPUT_SIZE = 608 106 | INPUT_LIST_FILE = './ped_list.txt' 107 | SAVE_PATH = './images_results/' 108 | 109 | 110 | ENGINE_FILE_PATH = './engines/engine_' + str(INPUT_SIZE) + '.trt' 111 | filenames = [] 112 | with open(INPUT_LIST_FILE, 'r') as l: 113 | lines = l.readlines() 114 | for line in lines: 115 | filename = line.strip() 116 | filenames.append(filename) 117 | input_resolution_yolov3_HW = (INPUT_SIZE, INPUT_SIZE) 118 | preprocessor = PreprocessYOLO(input_resolution_yolov3_HW) 119 | postprocessor_args = {"yolo_masks": [(3, 4, 5), (0, 1, 2)], 120 | "yolo_anchors": [(8,34), (14,60), (23,94), (39,149), (87,291), (187,472)], 121 | "obj_threshold": 0.1, 122 | "nms_threshold": 0.3, 123 | "yolo_input_resolution": input_resolution_yolov3_HW} 124 | postprocessor = PostprocessYOLO(**postprocessor_args) 125 | output_shapes = OUTPUT_SHAPES_DICT[str(INPUT_SIZE)] 126 | with get_engine(ENGINE_FILE_PATH) as engine, engine.create_execution_context() as context: 127 | inputs, outputs, bindings, stream = common.allocate_buffers(engine) 128 | for filename in filenames: 129 | image_raw, image = preprocessor.process(filename) 130 | shape_orig_WH = image_raw.size 131 | trt_outputs = [] 132 | 133 | # Do inference 134 | print('Running inference on image {}...'.format(filename)) 135 | inputs[0].host = image 136 | c_time = 0 137 | t1 = time.time() 138 | trt_outputs = common.do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream) 139 | t2 = time.time() 140 | c_time = t2-t1 141 | print('Inference takes time {}...'.format(c_time)) 142 | trt_outputs = [output.reshape(shape) for output, shape in zip(trt_outputs, output_shapes)] 143 | 144 | boxes, classes, scores = postprocessor.process(trt_outputs, (shape_orig_WH)) 145 | if len(boxes) != 0: 146 | obj_detected_img = draw_bboxes(image_raw, boxes, scores, classes, ALL_CATEGORIES) 147 | else: 148 | obj_detected_img = image_raw 149 | savename_0 = filename.split('/')[-1] 150 | savename = savename_0.split('.')[0] 151 | output_image_path = SAVE_PATH + savename + '_' + str(INPUT_SIZE) + '.png' 152 | obj_detected_img.save(output_image_path, 'PNG') 153 | print('Saved image with bounding boxes of detected objects to {}.'.format(output_image_path)) 154 | 155 | if __name__ == '__main__': 156 | main() 157 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | tensorrt=5.0.6.3 2 | numpy>=1.15.1 3 | pycuda>=2017.1.1 4 | Pillow>=5.2.0 5 | --------------------------------------------------------------------------------