onnx推理python代码

前提

pip install onnxruntime-gpu

代码

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "4"


import onnxruntime as ort
import numpy as np
import SimpleITK as sitk

import math
import time

# 
# python -m tf2onnx.convert --saved-model 20240122/1 --opset 13 --output test.onnx
# 

def normalize_img(img:np.ndarray, hu_min_val, hu_max_val)->np.ndarray:
	""" 归一化 """
	# min_val=-1000
	# max_val=600
	value_range = hu_max_val - hu_min_val
	norm_0_1 = (img - hu_min_val) / value_range
	img = np.clip(2*norm_0_1-1,-1,1)        
	return img



if __name__ == '__main__':
    testPatchPath = "/home/lhb_python_2/lung_datasets/dataset_128_32/val/JPEGImages/135_0013.nii.gz"
    image = sitk.ReadImage(testPatchPath)
    imageArray = sitk.GetArrayFromImage(image)
    print(f'imageArray.shape={imageArray.shape}')
    
	# preprocessing
    img = normalize_img(imageArray, -1000, 600)
    img = np.transpose(img, axes=[1,2,0])
    img = img[np.newaxis, ..., np.newaxis].astype('float32')

    onnxModelPath = "./test.onnx"
    sessOrt = ort.InferenceSession(
          onnxModelPath,
          providers=[
                # "CPUExecutionProvider",
                "CUDAExecutionProvider"
          ]
    )
    print(f' load onnx model successfully! '.center(60,"☆"))

	# predict
    preds = sessOrt.run(
          	output_names=["outputs"],
          	input_feed={"input_1": img}
			)
	# print(np.array(preds).squeeze().shape)	
    pred = np.array(preds).squeeze()
    preds = np.argmax(pred, axis=-1)
    preds = np.transpose(preds,axes=[2,0,1]).astype('uint8')

    outImage = sitk.GetImageFromArray(preds)
    sitk.WriteImage(outImage, "test.nii.gz")