Tensorrt-Sample-python-yolov3_onnx

#!/usr/bin/env python3
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

from __future__ import print_function

import os
import sys

import numpy as np
import tensorrt as trt
from data_processing import ALL_CATEGORIES, PostprocessYOLO, PreprocessYOLO
from PIL import ImageDraw

sys.path.insert(1, os.path.join(sys.path[0], ".."))
from downloader import getFilePath

from samples.python import common

TRT_LOGGER = trt.Logger()  # 创建一个Tensorrt日志器，用于记录推理过程中的日志


# 这个函数使用PIL.ImageDraw绘制矩形框，并显示检测到的物理类别和置信度
def draw_bboxes(
    image_raw, bboxes, confidences, categories, all_categories, bbox_color="blue"
):
    """Draw the bounding boxes on the original input image and return it.

    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')
    """

    draw = ImageDraw.Draw(image_raw)
    print(bboxes, confidences, categories)
    for box, score, category in zip(bboxes, confidences, categories):
        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))

        draw.rectangle(((left, top), (right, bottom)), outline=bbox_color)
        draw.text(
            (left, top - 12),
            "{0} {1:.2f}".format(all_categories[category], score),
            fill=bbox_color,
        )

    return image_raw


def get_engine(onnx_file_path, engine_file_path=""):
    """Attempts to load a serialized engine if available, otherwise builds a new TensorRT engine and saves it."""
    #接受两个参数，一个是onnx模型的文件路径，第二个是序列化的TensorRT引擎文件路径，用于保存或加载TensorRT引擎
    def build_engine():  #定义一个内部函数，用于从ONNX文件构建一个TensorRT引擎
        """Takes an ONNX file and creates a TensorRT engine to run inference with"""
        with (trt.Builder(TRT_LOGGER) as builder, # trt.Builder(TRT_LOGGER):创建一个TensorRT构建器对象(builder),用于构建网络
              builder.create_network(0) as network,  # 创建一个空的网络定义对象(network),其中0表示默认的网络标志
              builder.create_builder_config() as config,  #构建一个构建配置对象，用于设置tensorrt引擎构建过程中的参数
              trt.OnnxParser(network,TRT_LOGGER) as parser, # 创建一个onnx解析器，用于将onnx文件解析并转换为Tensorrt网络
              trt.Runtime(TRT_LOGGER) as runtime): # 创建一个Tensorrt运行对象，用于加载和运行引擎

            #设置工作区内存池的大小，用于构建引擎的临时存储，这里设置为256Mib
            config.set_memory_pool_limit(
                trt.MemoryPoolType.WORKSPACE, 1 << 28
            )  # 256MiB

            # Parse model file检查onnx文件是否存在
            if not os.path.exists(onnx_file_path):
                print(
                    "ONNX file {} not found, please run yolov3_to_onnx.py first to generate it.".format(
                        onnx_file_path
                    )
                )
                exit(0)
            print("Loading ONNX file from path {}...".format(onnx_file_path))
            # 打开onnx文件并使用parser.pase解析文件，如果解析失败，则输出所有错误信息并返回None
            with open(onnx_file_path, "rb") as model:
                print("Beginning ONNX file parsing")
                if not parser.parse(model.read()):
                    print("ERROR: Failed to parse the ONNX file.")
                    for error in range(parser.num_errors):
                        print(parser.get_error(error))
                    return None


            # The actual yolov3.onnx is generated with batch size 64. Reshape input to batch size 1
            # 重塑网络的输入形状
            network.get_input(0).shape = [1, 3, 608, 608]

            print("Completed parsing of ONNX file")
            print(
                "Building an engine from file {}; this may take a while...".format(
                    onnx_file_path
                )
            )

            plan = builder.build_serialized_network(network, config) #调用函数，根据网络和配置构建TensorRT引擎，并将其序列化为二进制格式(plan)

            engine = runtime.deserialize_cuda_engine(plan)#调用函数将序列化的引擎反序列化为可以使用的Tensorrt引擎对象(engine)
            print("Completed creating Engine")
            #保存引擎
            with open(engine_file_path, "wb") as f:
                f.write(plan)
            return engine

    if os.path.exists(engine_file_path): #判断是否存在序列化的引擎文件
        # If a serialized engine exists, use it instead of building an engine.
        print("Reading engine from file {}".format(engine_file_path))
        with open(engine_file_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
            return runtime.deserialize_cuda_engine(f.read())
    else:
        return build_engine()


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

    # Try to load a previously generated YOLOv3-608 network graph in ONNX format:
    onnx_file_path = "yolov3.onnx"   #定义两个路径
    engine_file_path = "yolov3.trt"
    # Download a dog image and save it to the following file path:
    input_image_path = getFilePath("samples/python/yolov3_onnx/dog.jpg")
    # Two-dimensional tuple with the target network's (spatial) input resolution in HW ordered
    input_resolution_yolov3_HW = (608, 608)  #网络输入分辨率
    # Create a pre-processor object by specifying the required input resolution for YOLOv3
    preprocessor = PreprocessYOLO(input_resolution_yolov3_HW)  #预处理对象，通过传入YOLO3所需要的输入分辨率来初始化
    # Load an image from the specified input path, and return it together with  a pre-processed version
    image_raw, image = preprocessor.process(input_image_path)  #加载并预处理图像
    # Store the shape of the original input image in WH format, we will need it for later
    shape_orig_WH = image_raw.size  #保存原始图像的宽高，以供后续的使用

    # Output shapes expected by the post-processor
    output_shapes = [(1, 255, 19, 19), (1, 255, 38, 38), (1, 255, 76, 76)]
    # Do inference with TensorRT
    trt_outputs = []

    with get_engine(  #进行推理
        onnx_file_path, engine_file_path
    ) as engine, engine.create_execution_context() as context:  #创建推理的执行上下文
        inputs, outputs, bindings, stream = common.allocate_buffers(engine)  #分配输入、输出、绑定和CUDA流的内存缓冲区，用于推理过程中传递数据
        # Do inference
        print("Running inference on image {}...".format(input_image_path))
        # Set host input to the image. The common.do_inference function will copy the input to the GPU before executing.
        inputs[0].host = image  #将预处理后的图像 image 作为输入，赋值给输入缓冲区
        #调用common.do_inference函数进行推理，该函数会将输入数据拷贝到GPU，并执行推理操作 ，最后将结果拷贝回主机内存
        trt_outputs = common.do_inference(
            context, #推理上下文
            engine=engine,
            bindings=bindings, #绑定输入输出的内存
            inputs=inputs,
            outputs=outputs, #用于存储输入和输出数据的缓冲区
            stream=stream,#CUDA流，用于异步操作
        )

    # Before doing post-processing, we need to reshape the outputs as the common.do_inference will give us flat arrays.
    #重塑推理输出，由于推理后的结果是扁平化的一维数组，需要根据模型输出的形状将其恢复为多维数组
    trt_outputs = [
        output.reshape(shape) for output, shape in zip(trt_outputs, output_shapes)
    ]

    postprocessor_args = {
        "yolo_masks": [   #yolo_mask yolo3使用的掩玛，用于不同特征图上的检测
            (6, 7, 8),
            (3, 4, 5),
            (0, 1, 2),
        ],  # A list of 3 three-dimensional tuples for the YOLO masks
        "yolo_anchors": [  #yolo3与定义的anchor大小，用于检测不同尺度的对象
            (10, 13),
            (16, 30),
            (33, 23),
            (30, 61),
            (62, 45),  # A list of 9 two-dimensional tuples for the YOLO anchors
            (59, 119),
            (116, 90),
            (156, 198),
            (373, 326),
        ],
        "obj_threshold": 0.6,  # Threshold for object coverage, float value between 0 and 1 #检测对象的置信度阈值
        "nms_threshold": 0.5,  # Threshold for non-max suppression algorithm, float value between 0 and 1 #非极大值抑制的阈值
        "yolo_input_resolution": input_resolution_yolov3_HW,
    }

    postprocessor = PostprocessYOLO(**postprocessor_args)  #使用定义的参数创建yolov3的后处理器postprocessor

    # Run the post-processing algorithms on the TensorRT outputs and get the bounding box details of detected objects
    boxes, classes, scores = postprocessor.process(trt_outputs, (shape_orig_WH))  #运行后处理算法，提取检测结果的边框(boxes)、类别(classes)和置信度(scores)
    # Draw the bounding boxes onto the original input image and save it as a PNG file
    obj_detected_img = draw_bboxes(image_raw, boxes, scores, classes, ALL_CATEGORIES)  #调用draw_bboxes函数在原始图像上绘制检测到的物理边框
    output_image_path = "dog_bboxes.png"  #保存包含边框的图像
    obj_detected_img.save(output_image_path, "PNG")  #
    print(
        "Saved image with bounding boxes of detected objects to {}.".format(
            output_image_path
        )
    )

    # Free host and device memory used for inputs and outputs
    common.free_buffers(inputs, outputs, stream)  #释放内存防止内存泄漏


if __name__ == "__main__":
    main()