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