模型预测结果显示程序如何转C++

[hwang12345]

想使用TensorRT在Vavier平台运行,目前实现了pth模型转trt模型,并使用C++ 加载模型进行预测,但是对于输出结果怎么整理有点迷惑,寻求指教. 对应python代码:

 out_torch = out_torch[:, ::-1, :]
            prob_torch = scipy.special.softmax(out_torch[:-1, :, :], axis=0)

            idx = np.arange(cfg.griding_num) + 1
            idx = idx.reshape(-1, 1, 1)

            loc = np.sum(prob_torch * idx, axis=0)
            out_torch = np.argmax(out_torch, axis=0)
            loc[out_torch == cfg.griding_num] = 0
            out_torch = loc


            h_inputs[0].host = trt_input
            t3 = time.time()
            trt_outputs = common.do_inference_v2(context, bindings=bindings, inputs=h_inputs, outputs=h_outputs,
                                                 stream=stream) # 输出结果为一维数组
            t4 = time.time()
            # print(trt_outputs)

            out_trt = trt_outputs[0].reshape((101, 56, 4))  # 将一维数组转为 （w+1）* sample_rows * 4 的tensor数据
            out_trt = out_trt[:, ::-1, :]
            prob_trt = scipy.special.softmax(out_trt[:-1, :, :], axis=0)

            idx = np.arange(100) + 1
            idx = idx.reshape(-1, 1, 1)

            loc = np.sum(prob_trt * idx, axis=0)
            out_trt = np.argmax(out_trt, axis=0)
            loc[out_trt == 100] = 0
            out_trt = loc

            for i in range(out_torch.shape[1]):
                if np.sum(out_torch[:, i] != 0) > 2:
                    for k in range(out_torch.shape[0]):
                        if out_torch[k, i] > 0:
                            ppp = (int(out_torch[k, i] * col_sample_w * img_w / 800) - 1,
                                   int(img_h * (row_anchor[cls_num_per_lane - 1 - k] / 288)) - 1)
                            cv2.circle(frame, ppp, 5, (0, 255, 0), -1)