I don’t know how to use tract-onnx. Could anyone help, please?

[atopx]

I have implemented the inference method in Python, but I don’t know how to use tract-onnx. Could anyone help, please?

python code:

import cv2
import numpy as np
import onnxruntime

class OnnxModel():
    def __init__(self, onnxpath):
        self.onnx_session = onnxruntime.InferenceSession(onnxpath)
        self.input_name = [node.name for node in self.onnx_session.get_inputs()]
        self.output_name = [node.name for node in self.onnx_session.get_outputs()]

    def get_input_feed(self, img_tensor):
        input_feed = {}
        for name in self.input_name:
            input_feed[name] = img_tensor
        return input_feed

    def inference(self, img_path):
        img = cv2.imread(img_path)
        or_img = cv2.resize(img, (640, 640))
        img = or_img[:, :, ::-1].transpose(2, 0, 1)
        img = img.astype(dtype=np.float32) / 255.0
        img = np.expand_dims(img, axis=0)
        input_feed = self.get_input_feed(img)
        return self.onnx_session.run(None, input_feed)[0]

    @staticmethod
    def nms(dets, thresh):
        x1 = dets[:, 0]
        y1 = dets[:, 1]
        x2 = dets[:, 2]
        y2 = dets[:, 3]
        areas = (y2 - y1 + 1) * (x2 - x1 + 1)
        scores = dets[:, 4]
        keep = []
        index = scores.argsort()[::-1]
        while index.size > 0:
            i = index[0]
            keep.append(i)
            x11 = np.maximum(x1[i], x1[index[1:]])
            y11 = np.maximum(y1[i], y1[index[1:]])
            x22 = np.minimum(x2[i], x2[index[1:]])
            y22 = np.minimum(y2[i], y2[index[1:]])
            w = np.maximum(0, x22 - x11 + 1)
            h = np.maximum(0, y22 - y11 + 1)
            overlaps = w * h
            ious = overlaps / (areas[i] + areas[index[1:]] - overlaps)
            idx = np.where(ious <= thresh)[0]
            index = index[idx + 1]
        return keep

    @staticmethod
    def xywh2xyxy(x):
        # [x, y, w, h] to [x1, y1, x2, y2]
        y = np.copy(x)
        y[:, 0] = x[:, 0] - x[:, 2] / 2
        y[:, 1] = x[:, 1] - x[:, 3] / 2
        y[:, 2] = x[:, 0] + x[:, 2] / 2
        y[:, 3] = x[:, 1] + x[:, 3] / 2
        return y


    def filter_box(self, org_box, conf_thres, iou_thres):
        org_box = np.squeeze(org_box)
        conf = org_box[..., 4] > conf_thres
        box = org_box[conf == True]
        cls_cinf = box[..., 5:]
        cls = []
        for i in range(len(cls_cinf)):
            cls.append(int(np.argmax(cls_cinf[i])))
        all_cls = list(set(cls))
        output = []
        for i in range(len(all_cls)):
            curr_cls = all_cls[i]
            curr_cls_box = []
            curr_out_box = []
            for j in range(len(cls)):
                if cls[j] == curr_cls:
                    box[j][5] = curr_cls
                    curr_cls_box.append(box[j][:6])
            curr_cls_box = np.array(curr_cls_box)
            curr_cls_box = self.xywh2xyxy(curr_cls_box)
            curr_out_box = self.nms(curr_cls_box, iou_thres)
            for k in curr_out_box:
                output.append(curr_cls_box[k])
        return np.array(output)

    def dector(self, filepath: str, threshold=0.7) -> list[list]:
        output = self.inference(filepath)
        # return [ [left, top, right, bottom, confidence, label]... ]
        return model.filter_box(output, threshold, threshold)

if __name__ == "__main__":
    model = OnnxModel('model.onnx')
    result = model.dector('source-new.png')
    print(result)

output:

[[564.80676    194.62836    630.537      260.2662       0.9500251
    0.        ]
 [151.31021    570.7323     215.9321     634.4846       0.9397042
    0.        ]
 [426.17404    257.70303    491.2546     322.25327      0.9520246
    1.        ]
 [150.29614      7.7866936  215.57364     72.568375     0.9440128
    1.        ]
 [289.3657      70.556816   354.4152     134.56284      0.9500934
    2.        ]
 [220.8186       8.242142   285.3549      72.05295      0.94878364
    2.        ]
 [288.67944      8.145216   353.52643     72.0813       0.9469146
    3.        ]
 [ 82.1537     508.6173     146.89296    572.2512       0.94805664
    4.        ]
 [221.10785    571.15234    284.69308    634.3749       0.9454925
    4.        ]
 [ 12.84111    507.75473     79.19514    572.7668       0.9484613
    5.        ]
 [288.7487     194.96596    354.43326    260.1801       0.950827
    6.        ]
 [427.1264     195.4267     491.9975     259.71866      0.95043087
    6.        ]
 [220.36661    445.72968    285.6228     510.09424      0.9486555
    6.        ]
 [150.884      509.20374    214.34729    571.4712       0.946614
    7.        ]
 [358.01358    195.98413    421.327      259.08005      0.9500728
    8.        ]
 [494.7517     195.38573    560.04675    259.41656      0.9384693
    8.        ]
 [287.5537     508.28732    353.73376    573.03235      0.95035625
   10.        ]
 [564.4168     507.7809     630.37286    572.78796      0.95458126
   12.        ]
 [288.50757    383.26044    354.2539     447.4588       0.95292425
   13.        ]
 [ 42.63672     42.064148   604.17267    614.4774       0.9601581
   14.        ]]