1.环境准备

pip install onnx
pip install onnxruntime-gpu
pip install PIL
pip install paddleseg

2.动态图转静态图

#动转静
!python /home/aistudio/PaddleSeg/export.py \
       --config /home/aistudio/PaddleSeg/configs/quick_start/UNet.yml \
       --model_path /home/aistudio/动态图模型/UNet.pdparams

3.静态图转onnxr

# 动转静.静转onnx
!paddle2onnx \
    --model_dir /home/aistudio/静态图模型/UNet \
    --model_filename model.pdmodel\
    --params_filename model.pdiparams \
    --save_file onnx模型/Unet.onnx \
    --opset_version 12

4.利用onnxruntime进行推理

（1）检测导出onnx模型是否正确

import onnx
# 我们可以使用异常处理的方法进行检验
try:
    # 当我们的模型不可用时，将会报出异常/home/aistudio/onnx模型/Unet++2.onnx
    model = onnx.load("/home/aistudio/onnx模型/Unet.onnx")
    onnx.checker.check_model(model)
except onnx.checker.ValidationError as e:
    print("The model is invalid: %s"%e)
else:
    # 模型可用时，将不会报出异常，并会输出“The model is valid!”
    print("The model is valid!")
print(onnx.helper.printable_graph(model.graph))

（2）颜色映射

import os

import cv2
import numpy as np
from PIL import Image as PILImage


def visualize(image, result, color_map, save_dir=None, weight=0.6):
    """
    Convert predict result to color image, and save added image.

    Args:
        image (str): The path of origin image.
        result (np.ndarray): The predict result of image.
        color_map (list): The color used to save the prediction results.
        save_dir (str): The directory for saving visual image. Default: None.
        weight (float): The image weight of visual image, and the result weight is (1 - weight). Default: 0.6

    Returns:
        vis_result (np.ndarray): If `save_dir` is None, return the visualized result.
    """

    color_map = [color_map[i:i + 3] for i in range(0, len(color_map), 3)]
    color_map = np.array(color_map).astype("uint8")
    # Use OpenCV LUT for color mapping
    c1 = cv2.LUT(result, color_map[:, 0])
    c2 = cv2.LUT(result, color_map[:, 1])
    c3 = cv2.LUT(result, color_map[:, 2])
    pseudo_img = np.dstack((c3, c2, c1))

    im = cv2.imread(image)
    vis_result = cv2.addWeighted(im, weight, pseudo_img, 1 - weight, 0)

    if save_dir is not None:
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)
        image_name = os.path.split(image)[-1]
        out_path = os.path.join(save_dir, image_name)
        cv2.imwrite(out_path, vis_result)
    else:
        return vis_result


def get_pseudo_color_map(pred, color_map=None):
    """
    Get the pseudo color image.

    Args:
        pred (numpy.ndarray): the origin predicted image.
        color_map (list, optional): the palette color map. Default: None,
            use paddleseg's default color map.

    Returns:
        (numpy.ndarray): the pseduo image.
    """
    pred_mask = PILImage.fromarray(pred.astype(np.uint8), mode='P')
    if color_map is None:
        color_map = get_color_map_list(256)
    pred_mask.putpalette(color_map)
    return pred_mask


def get_color_map_list(num_classes, custom_color=None):
    """
    Returns the color map for visualizing the segmentation mask,
    which can support arbitrary number of classes.

    Args:
        num_classes (int): Number of classes.
        custom_color (list, optional): Save images with a custom color map. Default: None, use paddleseg's default color map.

    Returns:
        (list). The color map.
    """

    num_classes += 1
    color_map = num_classes * [0, 0, 0]
    for i in range(0, num_classes):
        j = 0
        lab = i
        while lab:
            color_map[i * 3] |= (((lab >> 0) & 1) << (7 - j))
            color_map[i * 3 + 1] |= (((lab >> 1) & 1) << (7 - j))
            color_map[i * 3 + 2] |= (((lab >> 2) & 1) << (7 - j))
            j += 1
            lab >>= 3
    color_map = color_map[3:]

    if custom_color:
        color_map[:len(custom_color)] = custom_color
    return color_map


def paste_images(image_list):
    """
    Paste all image to a image.
    Args:
        image_list (List or Tuple): The images to be pasted and their size are the same.
    Returns:
        result_img (PIL.Image): The pasted image.
    """
    assert isinstance(image_list,
                      (list, tuple)), "image_list should be a list or tuple"
    assert len(
        image_list) > 1, "The length of image_list should be greater than 1"

    pil_img_list = []
    for img in image_list:
        if isinstance(img, str):
            assert os.path.exists(img), "The image is not existed: {}".format(
                img)
            img = PILImage.open(img)
            img = np.array(img)
        elif isinstance(img, np.ndarray):
            img = PILImage.fromarray(img)
        pil_img_list.append(img)

    sample_img = pil_img_list[0]
    size = sample_img.size
    for img in pil_img_list:
        assert size == img.size, "The image size in image_list should be the same"

    width, height = sample_img.size
    result_img = PILImage.new(sample_img.mode,
                              (width * len(pil_img_list), height))
    for i, img in enumerate(pil_img_list):
        result_img.paste(img, box=(width * i, 0))

    return result_img

（3）onnxruntime推理

%matplotlib inline
import time
import numpy as np
import onnxruntime
from onnxruntime import InferenceSession
import matplotlib.pyplot as plt
import cv2
# 加载ONNX模型
device_name = 'cuda:0' # or 'cpu'
# print(onnxruntime.get_available)

if device_name == 'cpu':
    providers = ['CPUExecutionProvider']
elif device_name == 'cuda:0':
    providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']

model = InferenceSession('/home/aistudio/onnx模型/Unet++.onnx',providers=providers)

# 准备输入
x=cv2.imread('newdata/origin/485.png')
x = cv2.cvtColor(x, cv2.COLOR_BGR2RGB)
x = cv2.resize(x, (512, 512))
mean = np.array([0.5, 0.5, 0.5]) * 255
std = np.array([0.5, 0.5, 0.5]) * 255
x = (x - mean) / std
print(x.max())
# 调整维度顺序从(H, W, C)到(C, H, W)
x = np.transpose(x, (2, 0, 1))

# 添加批处理维度，变成(1, C, H, W)
x = np.expand_dims(x, axis=0)
# x=np.tile(x,(2,1,1,1))
print('x',x.shape)
# 确保输入类型为float32
x = x.astype('float32')
input_name = model.get_inputs()[0].name
label_name = model.get_outputs()[0].name
# 模型预测
start = time.time()
# print('sess.get_inputs()[0].name',sess.get_inputs()[0].name)
ort_outs = model.run(None, input_feed={input_name: x})[0]

end = time.time()

print('ONNXRuntime predict time: %.04f s' % (end - start))
print('out',ort_outs.shape)

segmentation_mask = np.array(ort_outs)
print('segmentation_mask',segmentation_mask.shape)
segmentation_mask=segmentation_mask[0,:,:]
# [0,:,:,:]
# segmentation_mask=np.transpose(segmentation_mask, (1, 2, 0))
segmentation_mask1=get_pseudo_color_map(segmentation_mask)
print(segmentation_mask.max())
# 可视化分割掩码
# print('segmentation_mask',segmentation_mask1.shape)
print("像素类别",np.unique(segmentation_mask1))

plt.imshow(segmentation_mask1,cmap='jet') # 使用'jet'颜色映射来区分不同的类别
plt.colorbar() # 显示颜色条以理解颜色和类别的关系
plt.title("Segmentation Mask")
plt.show()