MMSegmentation笔记07：使用MMDeploy部署

1. 安装MMDeploy

• 首先要安装MMDeploy库，直接在对应环境下执行命令：pip install mmdeploy==1.2.0

• 然后在对应项目下安装MMDeploy源码：在cmd中执行：git clone https://github.com/open-mmlab/mmdeploy.git

• 然后在Python中验证是否安装成功：

import mmdeploy
print('MMDeploy 版本', mmdeploy.__version__)

2. 转换中间格式

有两种方式：

一种是直接用OpenMMLab提供的在线转换工具。

而是用源码中的tools/deploy.py文件在命令行中执行转换

python tools/deploy.py \
        configs/mmseg/segmentation_onnxruntime_dynamic.py \
        ../mmsegmentation/Zihao-Configs/ZihaoDataset_FastSCNN_20230818.py \
        ../mmsegmentation/checkpoint/Zihao_FastSCNN.pth \
        ../mmsegmentation/data/watermelon_test1.jpg \
        --work-dir mmseg2onnx_fastscnn \
        --dump-info

转换完成后，会在mmseg2onnx_fastscnn文件夹下生成一系列文件

• deploy.json：模型描述，用于MMDeploy Runtime推理

• detail.json：模型转ONNX的信息，用于追溯bug

• end2end.onnx：ONNX模型

• output_onnxruntime_0.png：使用推理框架预测的结果

• output_pytorch_0.png：使用pytorch预测的结果

• pipeline.json：模型输入、预处理、推理、后处理，每一步骤的输入输出信息，用于给使用者说明处理流程

3. 部署推理测试

1. 安装onnxruntime包

2. 图像预处理

3. 执行推理

本节的代码如下：

"""
==========================================
@author: Seaton
@Time: 2023/8/21:16:33
@IDE: PyCharm
@Summary：使用onnx格式推理
==========================================
"""
import numpy as np

import mmdeploy
import onnxruntime as ort
import cv2, torch
import matplotlib.pyplot as plt

# print('ONNXRuntime 版本', ort.__version__)
# print("MMDeploy 版本", mmdeploy.__version__)
onnx_path = 'mmdeploy/mmseg2onnx_fastscnn/end2end.onnx'
ort_session = ort.InferenceSession(onnx_path, provider=['CUDAExecutionProvider', 'CPUExecutionProvider'])

# 随机值预测
# x = torch.randn(1, 3, 1024, 2048).numpy()
# print(x.shape)
# ort_inputs = {'input': x}
# ort_output = ort_session.run(['output'], ort_inputs)[0]
# print(ort_output.shape)

# 真实图像预测
img_path = 'mmsegmentation/Watermelon87_Semantic_Seg_Mask/img_dir/val/dua-hau-1.jpg'
img_bgr = cv2.imread(img_path)
print(img_bgr.shape)
# 从原图中裁剪出宽高2:1的图像
h, w = img_bgr.shape[0], img_bgr.shape[1]
center_x, center_y = w // 2, h // 2
ratio = 3
new_h = 100 * ratio
new_w = 200 * ratio
img_bgr_crop = img_bgr[int(center_y-new_h/2):int(center_y+new_h/2), int(center_x - new_w/2):int(center_x + new_w/2)]
print(img_bgr_crop.shape)
img_bgr_resize = cv2.resize(img_bgr_crop, (2048, 1024))

# 预处理
img_tensor = img_bgr_resize
mean = (123.675, 116.28, 103.53)
std = (58.395, 57.12, 57.375)
img_tensor = (img_tensor - mean) / std
img_tensor = img_tensor.astype('float32')
img_tensor = cv2.cvtColor(img_tensor, cv2.COLOR_BGR2RGB)
img_tensor = np.transpose(img_tensor, (2, 0, 1))
input_tensor = np.expand_dims(img_tensor, axis=0)
print(input_tensor.shape)

# ONNX 预测
ort_inputs = {'input': input_tensor}
ort_output = ort_session.run(['output'], ort_inputs)[0]
pred_mask = ort_output[0][0]
# plt.imshow(pred_mask)
# plt.show()
palette = [
    ['background', [127, 127, 127]],
    ['red', [0, 0, 200]],
    ['green', [0, 200, 0]],
    ['white', [144, 238, 144]],
    ['seed-black', [30, 30, 30]],
    ['seed-white', [8, 189, 251]]
]

palette_dict = {}
for idx, each in enumerate(palette):
    palette_dict[idx] = each[1]

opacity = 0.3  # 透明度，越大越接近原图
# 将预测的整数ID，映射为对应类别的颜色
pred_mask_bgr = np.zeros((pred_mask.shape[0], pred_mask.shape[1], 3))
for idx in palette_dict.keys():
    pred_mask_bgr[np.where(pred_mask == idx)] = palette_dict[idx]
pred_mask_bgr = pred_mask_bgr.astype('uint8')
# 将语义分割预测图和原图叠加显示
img_bgr = cv2.addWeighted(img_bgr_resize, opacity, pred_mask_bgr, 1 - opacity, 0)
plt.imshow(img_bgr)
plt.show()