python 调用onnx,验证测试模型输出

# -*- coding: utf-8 -*- 
"""
Project: Eigenlanes
Creator: CHENRAN
Create time: 2022-06-28 11:43
IDE: PyCharm
Introduction:
"""
from PIL import Image

import onnxruntime
# import MNN
import numpy as np
import torchvision.transforms as transforms
import os
import time
import cv2

class Dataset_loader():
    def __init__(self):
        self.dir = dict()

        self.setting_for_image_param()
        self.setting_for_dataset_path()

        # image transform
        self.transform = transforms.Compose([transforms.Resize((self.height, self.width), interpolation=2),
                                             transforms.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.25, hue=0.1),
                                             transforms.ToTensor()])
        self.normalize = transforms.Normalize(mean=self.mean, std=self.std)
        self.reverse = transforms.Compose([transforms.Resize((self.org_height-self.crop_size, self.org_width), interpolation=2),
                                           ]) #transforms.ToTensor()

    def setting_for_image_param(self):
        self.org_height = 590
        self.org_width = 1640
        self.height = 320
        self.width = 800
        self.size = [self.width, self.height, self.width, self.height]
        self.mean = [0.485, 0.456, 0.406]
        self.std = [0.229, 0.224, 0.225]
        self.crop_size = 240 #240


    def get_image(self, idx):
        out = dict()
        img = Image.open(self.dir['dataset'] + '/{}'.format(idx)).convert('RGB')
        img.save(self.dir['out'] + '/' + idx[:-4]+'_final_ori.png')
        img = img.crop((0, self.crop_size, int(img.size[0]), int(img.size[1])))
        img = self.transform(img)
        out['img'] = self.normalize(img)
        return out

    def b_map_to_rgb_image(self, data):
        data_repeat = np.repeat(np.uint8(np.transpose(data, (1, 2, 0))*255), 3, 2)
        data = cv2.resize(data_repeat, (self.width, self.height))
        return data

    def padding_resize_to_originimg(self, idx,data):
        data = Image.fromarray(data)
        data_reverse = self.reverse(data)
        data_reverse = np.asarray(data_reverse)

        data_ori = np.zeros((self.org_height,self.org_width, 3))
        data_ori[self.crop_size:,:,:] = data_reverse
        cv2.imwrite(self.dir['out'] + '/'+idx[:-4]+ '_ori_size_seg.png', data_ori)
        return data_ori

    def setting_for_dataset_path(self):
        self.dir['dataset'] = '/home/ubuntu/CRTEST/Eigenlanes/data/images'
        self.dir['out'] = '/home/ubuntu/CRTEST/Eigenlanes/result'
        self.dir['model'] = '/home/ubuntu/CRTEST/Eigenlanes/Modeling/culane/code/model.onnx'
        




def main():
    Process_merge = Dataset_loader()
    path = Process_merge.dir['dataset']
    data_list = os.listdir(path)
    for i in data_list:
        x = Process_merge.get_image(i)
        x = np.array(x['img']).astype(np.float32)
        x = x[np.newaxis,:, :,:]
        t1 = time.time()
        fps = 0.0

        # onnx
        if_verify_onnx = True
        if if_verify_onnx:
            session = onnxruntime.InferenceSession(Process_merge.dir['model'])
            inputs = {session.get_inputs()[0].name: x}
            outs = session.run(None, inputs)
            outs = Process_merge.b_map_to_rgb_image(outs[0][0])
            print('seg saving：{}'.format(i))
            cv2.imwrite(Process_merge.dir['out']+'/'+i[:-4]+'_seg.png', outs)
            Process_merge.padding_resize_to_originimg(i,outs)
            fps = 1 / (time.time() - t1)
            print('test fps: {}'.format(fps))


if __name__ == '__main__':
    main()

Result:

 

 

Reference:

1. reference