前提
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")