如何解析trtexec输出的模型推理结果数据

最新推荐文章于 2025-10-22 10:51:38 发布

原创

最新推荐文章于 2025-10-22 10:51:38 发布 · 2.9k 阅读

3 ·

CC 4.0 BY-SA版权

文章标签：

#tensorrt #trtexec #efficientdet #onnx

该博客介绍了如何使用trtexec工具解析ONNX模型高效Det-D0的输出，并将结果保存到JSON文件中。内容涉及了对输出数据的后处理，包括非极大值抑制（NMS），以在图像上绘制识别结果的边界框。同时，提供了使用PyTorch和ONNXRuntime的代码示例，展示如何将模型输出转换为与trtexec相同格式的JSON文件。

还是以onnx为例，执行下面的命令解析onnx生成engine并推理输出结果到一个json文件:

./trtexec --onnx=efficientdet-d0-s.onnx --loadInputs='data':o4_clip-1_raw_data.bin --saveEngine=efficientdet-d0-s.engine  --exportOutput=trtexec-result.json

这个--exportOutput指定的json文件里保存的是network的输出节点的数据，对于我修改过的efficientdet网络的输出节点有三个：regression、anchors和classification，那么这个json文件的内容类似如下：

[
  { "name" : "regression"
  , "dimensions" : "1x49104x4"
  , "values" : [ 5.78506, 6.88422, -1.56519, -17.5148, -0.113892, 0.545003, -1.52597, -0.767865, 1.00629, 0.163376, 0.0242282, -1.89316, 0.187985, 0.499627, -0.414611, -2.14051, 0.371698, 0.478594, -0.0126426, 2.01423, 0.997112, 0.517545, 1.88847, -0.707338, 0.157562, 0.0627687, 0.10975, -0.430063, 0.537361, 0.670655, 0.428142, ...]},
  { "name" : "anchors"
  , "dimensions" : "1x49104x4"
  , "values" : [ -12, -12, 20, 20, -7.2, -18.4, 15.2, 26.4, -18.4, -7.2, 26.4, 15.2, -16.1587, -16.1587, 24.1587, 24.1587, -10.1111, -24.2222, 18.1111, 32.2222, -24.2222, -10.1111, 32.2222, 18.1111, -21.3984, -21.3984, 29.3984, 29.3984, -13.7789, -31.5578, 21.7789, 39.5578, -31.5578, -13.7789, 39.5578, 21.7789, -12, ...]},
  { "name" : "classification"
  , "dimensions" : "1x49104x1"
  , "values" : [ 0.000323705, 3.062e-07, 0.00457684, 0.000632869, 0.0004986, 0.000207652, 0.000125256, 6.19738e-07, 0.0203817, 8.39792e-06, 8.40121e-09, 9.50497e-05, 9.16859e-06, 2.48826e-05, 1.39859e-06, 3.46441e-06, 2.93581e-08, 0.000207522, 6.74278e-06, 1.8361e-08, 3.44744e-05, 3.35026e-06, 2.89377e-05, 3.85066e-07, 1.02452e-05,...]}
]

怎么把它这个结果解析出来并过滤处理成最终的结果并且在图上画出识别结果的bbox来呢？这个和一般的模型的post process是类似的，需要结合网络本身设计对输出数据做后处理得出合理的bbox和对应的class以及score，针对我这个修改过的efficientdet导出的onnx调用trtexec推理得到的json结果数据进行解析的相关代码如下:

#encoding: utf-8
import numpy as np
import os
import cv2
import json

classes = ['baggage']

def nms(bboxes, classifictaion, thresh):
    
    x1 = bboxes[:, 0]
    y1 = bboxes[:, 1]
    x2 = bboxes[:, 2]
    y2 = bboxes[:, 3]
    scores = classifictaion[:, 1]

    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    order = scores.argsort()[::-1]

    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.min