# -*-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]),
])