Conversion on ONNX inference code to C++

I have used YOLO v4 to train a custom model and converted the .pth into onnx and tensorrt model.
I want to write the inference API into C++ but I am not very good with C++. Tried searching for some forums but couldn’t find much help.
Could you please help me with this?

import argparse
import os
import cv2
import numpy as np
import onnxruntime
from src.model.yolox.data.datasets import COCO_CLASSES
from src.model.yolox.utils import mkdir, multiclass_nms, vis


class ONNXINFERENCE:

    def __init__(self):
        self.input_shape = "640,640"
        self.input_shape = tuple(map(int, self.input_shape.split(',')))
        self.image_path = "data/raw/image.png"
        self.origin_img = cv2.imread(self.image_path)
        self.mean = (0.485, 0.456, 0.406)
        self.std = (0.229, 0.224, 0.225)
        self.path = "model/yolox_custom.onnx"
        self.output_dir = "data/processed/"
        self.process()

    def process(self):
        self.main()

    def demo_postprocess(self, outputs, img_size, p6=False):
        grids = []
        expanded_strides = []

        if not p6:
            strides = [8, 16, 32]
        else:
            strides = [8, 16, 32, 64]

        hsizes = [img_size[0] // stride for stride in strides]
        wsizes = [img_size[1] // stride for stride in strides]

        for hsize, wsize, stride in zip(hsizes, wsizes, strides):
            xv, yv = np.meshgrid(np.arange(wsize), np.arange(hsize))
            grid = np.stack((xv, yv), 2).reshape(1, -1, 2)
            grids.append(grid)
            shape = grid.shape[:2]
            expanded_strides.append(np.full((*shape, 1), stride))

        grids = np.concatenate(grids, 1)
        expanded_strides = np.concatenate(expanded_strides, 1)
        outputs[..., :2] = (outputs[..., :2] + grids) * expanded_strides
        outputs[..., 2:4] = np.exp(outputs[..., 2:4]) * expanded_strides

        return outputs

    def preprocess(self, image, input_size, mean, std, swap=(2, 0, 1)):
        if len(image.shape) == 3:
            padded_img = np.ones((input_size[0], input_size[1], 3)) * 114.0
        else:
            padded_img = np.ones(input_size) * 114.0
        img = np.array(image)
        r = min(input_size[0] / img.shape[0], input_size[1] / img.shape[1])
        resized_img = cv2.resize(
            img,
            (int(img.shape[1] * r), int(img.shape[0] * r)),
            interpolation=cv2.INTER_LINEAR,
        ).astype(np.float32)
        padded_img[: int(img.shape[0] * r), : int(img.shape[1] * r)] = resized_img

        padded_img = padded_img[:, :, ::-1]
        padded_img /= 255.0
        if mean is not None:
            padded_img -= mean
        if std is not None:
            padded_img /= std
        padded_img = padded_img.transpose(swap)
        padded_img = np.ascontiguousarray(padded_img, dtype=np.float32)
        return padded_img, r

    def main(self):
        img, ratio = self.preprocess(self.origin_img, self.input_shape, self.mean, self.std)
        session = onnxruntime.InferenceSession("model/yolox_custom.onnx")
        ort_inputs = {session.get_inputs()[0].name: img[None, :, :, :]}
        output = session.run(None, ort_inputs)
        predictions = self.demo_postprocess(output[0], self.input_shape, p6=False)[0]
        boxes = predictions[:, :4]
        scores = predictions[:, 4:5] * predictions[:, 5:]
        boxes_xyxy = np.ones_like(boxes)
        boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2] / 2.
        boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3] / 2.
        boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2] / 2.
        boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.
        boxes_xyxy /= ratio
        score_thr = 0.7
        dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.45, score_thr=0.7)
        if dets is not None:
            final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5]
            self.origin_img = vis(self.origin_img, final_boxes, final_scores, final_cls_inds,
                                  conf=score_thr, class_names=COCO_CLASSES, )

        output_path = os.path.join(self.output_dir, self.image_path.split("/")[-1])
        cv2.imwrite(output_path, self.origin_img)


if __name__ == '__main__':
    call_onnx_inference = ONNXINFERENCE()
    call_onnx_inference.main()

    del call_onnx_inference

Moved to TensorRT forum

Hi,
Request you to share the ONNX model and the script if not shared already so that we can assist you better.
Alongside you can try few things:
https://docs.nvidia.com/deeplearning/tensorrt/quick-start-guide/index.html#onnx-export

  1. validating your model with the below snippet

check_model.py

import sys
import onnx
filename = yourONNXmodel
model = onnx.load(filename)
onnx.checker.check_model(model).
2) Try running your model with trtexec command.
https://github.com/NVIDIA/TensorRT/tree/master/samples/opensource/trtexec
In case you are still facing issue, request you to share the trtexec “”–verbose"" log for further debugging
Thanks!

Hi @NVES,

Thank you for the reply.

PFA the attached model.onnx.
yolox_custom.onnx (34.1 MB)

The model inference is running with the python code.
Just need help with C++ inference. I am unable to find any good solution for C++ inference code.

I am able to convert the model.onnx to tensorRT and that is working as well with Python inference API.

Could you please help me converting the Python API to C++ for the ONNX model?

Hi @NVES,

I pulled the TensorRT docker container from NGC and setup on my DGX. I updated the sample script for Yolox custom model. But it is giving me Segmentation fault. Attached the code script and error.log
sample.cpp.txt (11.6 KB)
error.log (109.3 KB)

Hi,

Script you have shared looks normal, loos like there is some problem with dimension related.
Please refer to following samples for better understanding.