本文为稀土掘金技术社区首发签约文章,30天内禁止转载,30天后未获授权禁止转载,侵权必究!
🍊作者简介:秃头小苏,致力于用最通俗的语言描述问题🍊专栏推荐:深度学习网络原理与实战
🍊近期目标:写好专栏的每一篇文章
🍊支持小苏:点赞👍🏼、收藏⭐、留言📩
写在前面
在前面的几节中,我已经为大家介绍了深度学习在分类模型中的部署,分为以下三节,不清楚的可以去看一看:
- 深度学习模型部署篇------从0部署深度学习分类模型(一)🍁🍁🍁
- 深度学习模型部署篇------从0部署深度学习分类模型(二)🍁🍁🍁
- 深度学习模型部署篇------利用Flask实现深度学习模型部署(三)🍁🍁🍁
那么这篇我将和大家唠唠如何部署语义分割模型,大家千万不要潜意识里觉得语义分割模型很难,其实它是很简单的,不清楚的可以去看看语义分割的开山之作------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,你会发现极其极其低。
小结
本篇就为大家介绍到这里了喔,希望大家都有所收获。⛳⛳⛳吐槽一下,今天踢球又一次伤到脚踝了,坏了,又得在宿舍呆好多天,这篇后面感觉自己写的有点仓促了,大家多担待,不说了,上床躺尸。🛌🏽🛌🏽🛌🏽
如若文章对你有所帮助,那就🛴🛴🛴