DATAWHALE 第45期组队学习__图像分类模型部署

---

一、部署ImageNet预训练图像分类模型

只能在本地跑代码，云平台无法跑起来

1. 导出ONNX模型：

导入工具包

Pytorch模型转ONNX模型

然后使用Netron对onnx模型可视化

二、推理引擎ONNX Runtime部署-预测摄像头实时画面

1.导入工具包

2.载入 onnx 模型，获取 ONNX Runtime 推理器

ort_session = onnxruntime.InferenceSession('resnet18.onnx')

3.载入ImageNet 1000图像分类标签和图像预处理

# 载入ImageNet 1000图像分类标签
df = pd.read_csv('imagenet_class_index.csv')
idx_to_labels = {}
for idx, row in df.iterrows():
    idx_to_labels[row['ID']] = row['Chinese']
    
from torchvision import transforms
# 测试集图像预处理-RCTN：缩放裁剪、转 Tensor、归一化
test_transform = transforms.Compose([transforms.Resize(256),
                                     transforms.CenterCrop(256),
                                     transforms.ToTensor(),
                                     transforms.Normalize(
                                         mean=[0.485, 0.456, 0.406], 
                                         std=[0.229, 0.224, 0.225])
                                    ])

4.获取摄像头的一帧画面

# 导入opencv-python
import cv2
import time
# 获取摄像头，传入0表示获取系统默认摄像头
cap = cv2.VideoCapture(1)
# 打开cap
cap.open(0)
time.sleep(1)
success, img_bgr = cap.read()    
# 关闭摄像头
cap.release()
# 关闭图像窗口
cv2.destroyAllWindows()

img_rgb=cv2.cvtColor(img_bgr,cv2.COLOR_BGR2RGB)
img_pil=Image.fromarray(img_rgb)
img_pil

5.将预测结果可视化在拍摄的图片上

# 预处理
input_img = test_transform(img_pil)
input_tensor = input_img.unsqueeze(0).numpy()
# onnx runtime 预测

# onnx runtime 输入
ort_inputs = {'input': input_tensor}
# onnx runtime 输出
pred_logits = ort_session.run(['output'], ort_inputs)[0]
pred_logits = torch.tensor(pred_logits)
import torch.nn.functional as F
pred_softmax = F.softmax(pred_logits, dim=1) # 对 logit 分数做 softmax 运算
## 解析图像分类预测结果
n = 5
top_n = torch.topk(pred_softmax, n) # 取置信度最大的 n 个结果
pred_ids = top_n[1].cpu().detach().numpy().squeeze() # 解析出类别
confs = top_n[0].cpu().detach().numpy().squeeze() # 解析出置信度
draw = ImageDraw.Draw(img_pil) 
# 在图像上写字
for i in range(len(confs)):
    pred_class = idx_to_labels[pred_ids[i]]
    text = '{:<15} {:>.3f}'.format(pred_class, confs[i])
    # 文字坐标，中文字符串，字体，rgba颜色
    draw.text((50, 100 + 50 * i), text, font=font, fill=(255, 0, 0, 1))
img = np.array(img_pil) # PIL 转 array

plt.imshow(img)
plt.show()

总结

ONNXRuntime是微软推出的一款推理框架，用户可以非常便利的用其运行一个onnx模型。从pytorch转到onnx，通过使用Netron对onnx模型可视化，可以非常清晰的理解深度学习框架的一些核心功能原理。