使用onnx推理语义分割

# -*-coding: utf-8 -*-
import numpy as np
import cv2
import onnxruntime

# 图片预处理
def image_pre(image):
    # BGR转为RGB
    img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    # 640*640这里填写你要resize大小
    img = cv2.resize(np.array(img), (640, 640), interpolation=cv2.INTER_CUBIC).astype(np.float32)
    img /= 255.0 
    img -= [0.485, 0.456, 0.406]
    img /= [0.229, 0.224, 0.225]
    #HWC--->CHW
    img = img.transpose(2, 0, 1)
    #CHW---> 1,C,H,W
    img = np.expand_dims(img, axis=0)
    return img

def coverLabelToGroup6(seg):
    #注意：这里label是RGB的,pattle填写你的类别对应的label标签，
    palette = [[109, 63, 21], [0, 102, 0], [0, 255, 0], [255, 153, 203], [255, 255, 0],
               [255, 0, 0], [101, 0, 0], [0, 0, 255], [154, 205, 255], [127, 127, 127]]
    color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)
    for label, color in enumerate(palette):
        color_seg[seg == label, :] = color
    return color_seg

# 输入图片的路径
image = cv2.imread("")
img = image_pre(image)
# onnx路径
onnx_model_path = ""
resnet_session = onnxruntime.InferenceSession(onnx_model_path)
inputs = {resnet_session.get_inputs()[0].name: img}
outs = resnet_session.run(None, inputs)

real_outs = outs[0].argmax(axis=1)[0]
output = coverLabelToGroup6(real_outs)
output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
cv2.imshow("output",output)
cv2.waitKey(0)

因为不想下载torch和torchvision，直接转onnx推理了，直接用cv2代替

def get_test_transform():
    return transforms.Compose([
        transforms.Resize([640, 640]),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])