深度学习模型部署篇——从0部署深度学习语义分割模型

本文为稀土掘金技术社区首发签约文章，30天内禁止转载，30天后未获授权禁止转载，侵权必究！

🍊作者简介：深度学习模型部署篇——从0部署深度学习分类模型（一）🍁🍁🍁

深度学习模型部署篇——从0部署深度学习分类模型（二）🍁🍁🍁

深度学习模型部署篇——利用Flask实现深度学习模型部署（三）🍁🍁🍁

那么这篇我将和大家唠唠如何部署语义分割模型，大家千万不要潜意识里觉得语义分割模型很难，其实它是很简单的，不清楚的可以去看看语义分割的开山之作——FCN，我也做过原理详解和源码实战篇，不清楚的可以去看看，相信你定会有所收获：

• 深度学习语义分割篇——FCN原理详解篇🌲🌲🌲
• 深度学习语义分割篇——FCN源码解析篇🌲🌲🌲

知道了语义分割的基本知识，你就可以来看这篇博客啦。当然了，我还做过语义分割的其它系列，感兴趣的欢迎去踩我的主页。🥂🥂🥂

准备好了喵，我们这就发车~~~🚖🚖🚖

这里我想给大家说一下本文的行文安排，首先我会基于最基础的语义分割模型FCN，给大家介绍语义分割模型部署的流程，熟悉这个流程之后为大家介绍介绍mmcv库，基于这个库来实现语义分割模型部署。🍄🍄🍄

基于FCN部署语义分割模型

好啦，我们这就开始了喔，本小节使用的是FCN模型进行模型部署，因此建议你先读读我写在前面中给出的两篇文章，对FCN实现语义分割有一个基本的了解。🥗🥗🥗

导入工具包

 import onnxruntime
 import numpy as np
 import cv2
 from src import fcn_resnet50
 import torch
 ​
 import matplotlib.pyplot as plt
 %matplotlib inline

大家注意一下这里from src import fcn_resnet50，我们从src下的fcn_model.py文件中导入了fcn模型结构，src下面的目录结构是这样的：

大家这里可能有点疑惑，从fcn_model.py中导入文件，为什么这里只写了from src呢，其实其它的内容写在了__init__.py文件中，文件内容如下：

对于此不熟悉的可以看一下我的这篇文章。🍡🍡🍡

创建模型

 aux = False
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 classes = 20
 # create model
 model = fcn_resnet50(aux=aux, num_classes=classes+1)
 ​
 weights_path = "./model_29.pth"
 # delete weights about aux_classifier
 weights_dict = torch.load(weights_path, map_location='cpu')['model']
 for k in list(weights_dict.keys()):
     if "aux" in k:
         del weights_dict[k]

这部分其实就是FCN预测部分的代码，我直接复制过来了。🍮🍮🍮

加载模型并设置为推理模式

 # 加载模型
 model.load_state_dict(weights_dict)
 # 设置为推理模式
 model = model.eval().to(device)

pytorch转ONNX格式

 x = torch.randn(1, 3, 520, 520).to(device)
 output = model(x)
 with torch.no_grad():
     torch.onnx.export(
         model,                   # 要转换的模型
         x,                       # 模型的任意一组输入
         'model_29.onnx', # 导出的 ONNX 文件名
         opset_version=11,        # ONNX 算子集版本
         input_names=['input'],   # 输入 Tensor 的名称（自己起名字）
         output_names=['output']  # 输出 Tensor 的名称（自己起名字）
     ) 
 ​

这部分和图像分类是一致的，需要注意的是这里的输入需要是（1, 3, 520, 520）的大小。

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

 # ONNX 模型路径
 onnx_path = './model_29.onnx'
 ort_session = onnxruntime.InferenceSession(onnx_path, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])

providers=['CUDAExecutionProvider', 'CPUExecutionProvider']参数表示自己根据硬件选择GPU或者CPU。

载入推理图像

 img_path = 'cat.jpg'
 img_bgr = cv2.imread(img_path)

我们来看看我们的小猫咪长什么样：

图像预处理

上面说到我们的输入需要是（1, 3, 520, 520）大小的，所以我们要进行resize操作：

 img_bgr_resize = cv2.resize(img_bgr, (520, 520)) # 缩放尺寸

enmmm，再来看看resize后的图像吧：【这是用plt库显示的结果】

接着还需要进行其它的一下预处理操作，如下：

 img_tensor = img_bgr_resize
 ​
 # BGR 三通道的均值
 mean = (123.675, 116.28, 103.53)
 ​
 # BGR 三通道的标准差
 std = (58.395, 57.12, 57.375)
 ​
 # 归一化
 img_tensor = (img_tensor - mean) / std
 img_tensor = img_tensor.astype('float32')
 ​
 # BGR 转 RGB
 img_tensor = cv2.cvtColor(img_tensor, cv2.COLOR_BGR2RGB)
 ​
 # 调整维度
 img_tensor = np.transpose(img_tensor, (2, 0, 1))   #h w c --> c h w
 # 扩充 batch-size 维度
 input_tensor = np.expand_dims(img_tensor, axis=0)

ONNX Runtime预测

 # ONNX Runtime 输入
 ort_inputs = {'input': input_tensor}
 # onnx runtime 输出
 ort_output = ort_session.run(['output'], ort_inputs)[0]

我们可以来看看ort_output的维度：

还记得我们要对这个维度怎么操作吗，如下：

我们要获取每个chanel中的最大值：

 pred_mask = ort_output.argmax(1)[0]  #获取ort_output数组中的第一个二维元素

此时pred_mask维度为：

我们可以通过np.unique(pred_mask)来看看pred_mask中有哪些值：

从上图可以看出有0和8，0表示背景，8表示cat。这个和voc类别文档中的是一致的，文档中背景没有写。

接着我们可以简单的实现一个可视化，先定义一个字典：

 # [127, 127, 127]表示灰色 ；[0, 180, 180]表示黄色
 palette_dict = {0: [127, 127, 127], 8: [0, 180, 180]}  

然后将0和8两个类别映射成不同的颜色：

 opacity = 0.2 # 透明度，越大越接近原图
 # 将预测的整数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')
 ​
 # 将语义分割预测图和原图叠加显示
 pred_viz = cv2.addWeighted(img_bgr_resize, opacity, pred_mask_bgr, 1-opacity, 0)

我们可以来看下最终的预测结果，即pred_viz，如下：

结果还是非常不错的。🥤🥤🥤

基于MMCV库实现语义分割模型部署

首先来介绍一下MMCV库，其是一个用于计算机视觉和多媒体计算的开源工具包，主要用于深度学习项目的开发和研究。MMCV是由中国科学院自动化研究所（Institute of Automation, Chinese Academy of Sciences）开发和维护的，它提供了许多用于图像和视频处理、计算机视觉任务、模型训练和部署的实用工具和组件。使用这个库我们可以很方便的实现各种计算机视觉任务，首先我们先来安装一下这个库：

 pip install -U openmim
 mim install mmengine
 mim install mmcv==2.0.0

接着安装其它的工具包：

 pip install opencv-python pillow matplotlib seaborn tqdm pytorch-lightning 'mmdet>=3.1.0' -i https://pypi.tuna.tsinghua.edu.cn/simple

这里使用了清华园镜像进行安装，会快很多。大家注意安装包的时候不要打开系统代理，不然会安装失败。由于我们要进行的是语义分割任务，因此需要下载mmsegmentation 源代码:

 git clone https://github.com/open-mmlab/mmsegmentation.git -b v1.1.1

然后我们进入mmsegmentation 目录安装MMSegmentation库:

 # 进入主目录
 import os
 os.chdir('mmsegmentation')
 # 安装`MMSegmentation`库
 pip install -v -e .

安装好后，我们来验证下我们是否安装成功：

 # 检查 mmcv
 import mmcv
 from mmcv.ops import get_compiling_cuda_version, get_compiler_version
 print('MMCV版本', mmcv.__version__)
 print('CUDA版本', get_compiling_cuda_version())
 print('编译器版本', get_compiler_version())

 # 检查 mmsegmentation
 import mmseg
 from mmseg.utils import register_all_modules
 from mmseg.apis import inference_model, init_model
 print('mmsegmentation版本', mmseg.__version__)

我的版本是这样的：

这些准备好了之后，我们就可以使用mmsegmentation进行语义分割任务了，操作也很简单，主要就是对各种配置文件的修改。由于本节主要介绍模型部署，这里如何进行训练就不叙述了。不清楚的可以点击☞☞☞了解详情。

这里我直接拿同济子豪兄得到的西瓜语义分割ONNX模型来举例了，呜呜，懒的自己从头训练了。🍋🍋🍋下载连接如下：

链接: pan.baidu.com/s/1E7Q0P79n… 提取码: luq8

下载完成后解压就得到了ONNX格式的模型。

这里先来简单介绍一下西瓜语义分割数据集：

一共有6个类别，分别为：

• 0：背景
• 1：红壤
• 2：绿壳
• 3：白皮
• 4：黒籽
• 5：白籽

下面是标注的图像：

我们有了ONNX模型后，我们就可以使用ONNX Runtime推理器进行推理了。

 # ONNX 模型路径
 onnx_path = 'mmseg2onnx_fastscnn/end2end.onnx'
 ort_session = onnxruntime.InferenceSession(onnx_path, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])

其实到这里就和前面基于FCN进行语义分割差不多了，我讲这部分就是想让大家了解了解MMCV这个工具，但是这节也没有详细介绍啦，后期可能会出一些基于MMCV实现目标检测、语义分割的博客。

西瓜数据集一个有6个类，它们的值都很接近，所以颜色很类似，我们将这6个值映射到不同的颜色，这样可视化出来更加美观：

 # 各类别的配色方案（BGR）
 palette = [     ['background', [127,127,127]],
     ['red', [0,0,200]],
     ['green', [0,200,0]],
     ['white', [144,238,144]],
     ['seed-black', [30,30,30]],
     ['seed-white', [180,180,180]]
 ]
 ​
 palette_dict = {}
 for idx, each in enumerate(palette):
     palette_dict[idx] = each[1]

这个代码不知道大家能否理解，有不理解的其实调试一下就会豁然开朗：

看了上图的调试过程，大家是不是就清晰了呢，最后palette_dict是一个字典，内容如下：

接下来，我们封装一个处理单张图像的函数，如下：【其实这部分就是把上节的几个小部分整合到了一起，没什么难度】

 opacity = 0.2 # 透明度，越大越接近原图
 def process_frame(img_bgr):
     
     '''
     输入摄像头画面 bgr-array，输出图像 bgr-array
     '''
     
     # 记录该帧开始处理的时间
     start_time = time.time()
     
     # 从原图中裁剪出高宽比1:2的最大图像
     h, w = img_bgr.shape[0], img_bgr.shape[1]
     new_h = w // 2 # 横屏图片，截取一半的宽度，作为新的高度
     img_bgr_crop = img_bgr[0:new_h, :]
     
     # 缩放至模型要求的高1024 x 宽2048像素
     img_bgr_resize = cv2.resize(img_bgr_crop, (2048, 1024)) # 缩放尺寸
     
     # 预处理
     img_tensor = img_bgr_resize
     mean = (123.675, 116.28, 103.53) # BGR 三通道的均值
     std = (58.395, 57.12, 57.375) # BGR 三通道的标准差
 ​
     # 归一化
     img_tensor = (img_tensor - mean) / std
     img_tensor = img_tensor.astype('float32')
     img_tensor = cv2.cvtColor(img_tensor, cv2.COLOR_BGR2RGB) # BGR 转 RGB
     img_tensor = np.transpose(img_tensor, (2, 0, 1)) # 调整维度
     input_tensor = np.expand_dims(img_tensor, axis=0) # 扩充 batch-size 维度
 ​
     # ONNX Runtime预测
     # ONNX Runtime 输入
     ort_inputs = {'input': input_tensor}
     # onnx runtime 输出
     ort_output = ort_session.run(['output'], ort_inputs)[0]
     pred_mask = ort_output[0][0]
 ​
     # 将预测的整数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')
 ​
     # 将语义分割预测图和原图叠加显示
     pred_viz = cv2.addWeighted(img_bgr_resize, opacity, pred_mask_bgr, 1-opacity, 0)
     
     img_bgr = pred_viz
     
     # 记录该帧处理完毕的时间
     end_time = time.time()
     # 计算每秒处理图像帧数FPS
     FPS = 1/(end_time - start_time)
 ​
     # 在画面上写字：图片，字符串，左上角坐标，字体，字体大小，颜色，字体粗细
     scaler = 2 # 文字大小
     FPS_string = 'FPS {:.2f}'.format(FPS) # 写在画面上的字符串
     img_bgr = cv2.putText(img_bgr, FPS_string, (25 * scaler, 100 * scaler), cv2.FONT_HERSHEY_SIMPLEX, 1.25 * scaler, (255, 0, 255), 2 * scaler)
     
     return img_bgr

接着我们可以来调用摄像头一帧图像来看看结果：

 # 获取摄像头，0为电脑默认摄像头，1为外接摄像头
 cap = cv2.VideoCapture(0)
 ​
 # 拍照
 time.sleep(1) # 运行本代码后等几秒拍照
 # 从摄像头捕获一帧画面
 success, frame = cap.read()
 ​
 cap.release() # 关闭摄像头
 cv2.destroyAllWindows() # 关闭图像窗口

frame的结果如下：

enmmmm，我没有西瓜啊，所以这个是手机里面的图片😂😂😂

接着使用img_bgr = process_frame(frame)来推理这张图片，分割结果如下：

可以看到，FPS为1.71，针对我的电脑它已经很快了，大家可以试试不使用ONNX Runtime推理的FPS，你会发现极其极其低。

小结

本篇就为大家介绍到这里了喔，希望大家都有所收获。⛳⛳⛳吐槽一下，今天踢球又一次伤到脚踝了，坏了，又得在宿舍呆好多天，这篇后面感觉自己写的有点仓促了，大家多担待，不说了，上床躺尸。🛌🏽🛌🏽🛌🏽

 
 

如若文章对你有所帮助，那就🛴🛴🛴