当前位置:首页> AI教程> 深度学习模型部署篇——从0部署深度学习语义分割模型

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

释放双眼,带上耳机,听听看~!
本文介绍了使用FCN模型进行语义分割模型部署的流程,包括模型创建、pytorch转ONNX格式、载入ONNX模型和ONNX Runtime预测等步骤。同时还介绍了图像预处理等相关操作。

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

🍊作者简介:深度学习模型部署篇——从0部署深度学习分类模型(一)🍁🍁🍁

  • 深度学习模型部署篇——从0部署深度学习分类模型(二)🍁🍁🍁
  • 深度学习模型部署篇——利用Flask实现深度学习模型部署(三)🍁🍁🍁
  • 那么这篇我将和大家唠唠如何部署语义分割模型,大家千万不要潜意识里觉得语义分割模型很难,其实它是很简单的,不清楚的可以去看看语义分割的开山之作——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下面的目录结构是这样的:

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

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

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

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

    创建模型

     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)
    

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

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

    图像预处理

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

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

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

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

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

     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的维度:

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

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

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

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

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

    此时pred_mask维度为:

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

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

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

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

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

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

     # [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,如下:

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

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

    基于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__)
    

    我的版本是这样的:

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

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

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

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

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

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

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

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

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

    一共有6个类别,分别为:

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

    下面是标注的图像:

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

    我们有了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]
    

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

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

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

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

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

     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的结果如下:

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

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

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

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

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

    小结

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

     
     

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

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

    本网站的内容主要来自互联网上的各种资源,仅供参考和信息分享之用,不代表本网站拥有相关版权或知识产权。如您认为内容侵犯您的权益,请联系我们,我们将尽快采取行动,包括删除或更正。
    AI教程

    深度学习在工业缺陷检测中的应用详解

    2023-11-19 10:18:14

    AI教程

    AI抠图使用指南:Stable Diffusion WebUI Rembg实用技巧

    2023-11-19 10:20:14

    个人中心
    购物车
    优惠劵
    今日签到
    有新私信 私信列表
    搜索