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

本文为稀土掘金技术社区首发签约文章,30天内禁止转载,30天后未获授权禁止转载,侵权必究!
🍊作者简介:秃头小苏,致力于用最通俗的语言描述问题

🍊专栏推荐:深度学习网络原理与实战

🍊近期目标:写好专栏的每一篇文章

🍊支持小苏:点赞👍🏼、收藏⭐、留言📩

写在前面

在前面的几节中,我已经为大家介绍了深度学习在分类模型中的部署,分为以下三节,不清楚的可以去看一看:

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

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

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


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


基于FCN部署语义分割模型

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

导入工具包

javascript 复制代码
 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文件中,文件内容如下:

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

创建模型

ini 复制代码
 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预测部分的代码,我直接复制过来了。🍮🍮🍮

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

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

pytorch转ONNX格式

ini 复制代码
 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推理器

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

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

载入推理图像

ini 复制代码
 img_path = 'cat.jpg'
 img_bgr = cv2.imread(img_path)

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

图像预处理

上面说到我们的输入需要是(1, 3, 520, 520)大小的,所以我们要进行resize操作:

ini 复制代码
 img_bgr_resize = cv2.resize(img_bgr, (520, 520)) # 缩放尺寸

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

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

ini 复制代码
 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预测

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

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

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

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

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

此时pred_mask维度为:

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

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

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

css 复制代码
 # [127, 127, 127]表示灰色 ;[0, 180, 180]表示黄色
 palette_dict = {0: [127, 127, 127], 8: [0, 180, 180]}  

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

ini 复制代码
 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)开发和维护的,它提供了许多用于图像和视频处理、计算机视觉任务、模型训练和部署的实用工具和组件。使用这个库我们可以很方便的实现各种计算机视觉任务,首先我们先来安装一下这个库:

ini 复制代码
 pip install -U openmim
 mim install mmengine
 mim install mmcv==2.0.0

接着安装其它的工具包:

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

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

bash 复制代码
 git clone https://github.com/open-mmlab/mmsegmentation.git -b v1.1.1

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

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

安装好后,我们来验证下我们是否安装成功:

python 复制代码
 # 检查 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())
python 复制代码
 # 检查 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推理器进行推理了。

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

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

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

less 复制代码
 # 各类别的配色方案(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是一个字典,内容如下:

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

ini 复制代码
 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

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

ini 复制代码
 # 获取摄像头,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,你会发现极其极其低。

小结

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

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

相关推荐
攻城狮_Dream3 分钟前
“探索未来医疗:生成式人工智能在医疗领域的革命性应用“
人工智能·设计·医疗·毕业
学习前端的小z32 分钟前
【AIGC】如何通过ChatGPT轻松制作个性化GPTs应用
人工智能·chatgpt·aigc
埃菲尔铁塔_CV算法1 小时前
人工智能图像算法:开启视觉新时代的钥匙
人工智能·算法
EasyCVR1 小时前
EHOME视频平台EasyCVR视频融合平台使用OBS进行RTMP推流,WebRTC播放出现抖动、卡顿如何解决?
人工智能·算法·ffmpeg·音视频·webrtc·监控视频接入
打羽毛球吗️1 小时前
机器学习中的两种主要思路:数据驱动与模型驱动
人工智能·机器学习
好喜欢吃红柚子1 小时前
万字长文解读空间、通道注意力机制机制和超详细代码逐行分析(SE,CBAM,SGE,CA,ECA,TA)
人工智能·pytorch·python·计算机视觉·cnn
小馒头学python1 小时前
机器学习是什么?AIGC又是什么?机器学习与AIGC未来科技的双引擎
人工智能·python·机器学习
神奇夜光杯2 小时前
Python酷库之旅-第三方库Pandas(202)
开发语言·人工智能·python·excel·pandas·标准库及第三方库·学习与成长
正义的彬彬侠2 小时前
《XGBoost算法的原理推导》12-14决策树复杂度的正则化项 公式解析
人工智能·决策树·机器学习·集成学习·boosting·xgboost
Debroon2 小时前
RuleAlign 规则对齐框架:将医生的诊断规则形式化并注入模型,无需额外人工标注的自动对齐方法
人工智能