语义分割——mmsegmentation框架使用

tang_seven

已于 2024-06-11 17:06:20 修改

阅读量1.8k

点赞数 20

文章标签：人工智能计算机视觉

于 2024-06-11 16:53:39 首次发布

本文链接：https://blog.csdn.net/tang_seven/article/details/139506727

版权

1.mmsegmentation简介

2.mmsegmentation安装

3.mmsegmentation使用（代码结构介绍）

4.mmsegmentation使用实战（deeplab v3+为例）

4.1配置

4.2训练：

4.3预测：

1.mmsegmentation简介

mmsegmentation 是 OpenMMLab 开源的基于 PyTorch 实现的功能强大的语义分割工具箱。

简单来说，就是一个集成了并实现了多个语义分割框架的代码框架，你可以在里面自由地选用对应的语义分割模型，比如FCN、U-net、PSPNet、Deeplab v3、Swim Transformer等，而无需自己编写对应的模型代码，只需要按照MMSegmentation 官网文档编写其对应的训练文件。甚至你可以直接下载其在某些数据上预训练后的权重文件，并进行预测，获得语义分割结果。

同时，mmsegmentation 还包含11 种主干网络代码，代表着你可以使用它进行自由组合，设计出你想要的语义分割网络。

所以，mmsegmentation 可以帮我们做什么呢？

如果你需要一些公开数据集（如voc，coco等）语义分割模型的实际比对结果，你就可以使用mmsegmentation 快速地获取结果，避免了一个一个模型进行手敲，并进行训练，节省大量时间。

再者，编写根据既定的规则编写训练以及预测代码总比编写模型代码简单的多。

2.mmsegmentation安装

首先，你得知道mmsegmentation 是基于pytorch实现的，所以你需要先安装pytorch，如果你需要使用GPU进行训练等，你还需要安装cuda等，具体网上都有教程，此处不进行叙述。

然后，你需要在安装好pytorch的环境中安装mmsegmentation 对应的依赖：

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

然后你可以直接下载mmsegmentation 的源码进行使用了：

https://gitcode.com/open-mmlab/mmsegmentation

或者通过命令行形式：

git clone -b main https://github.com/open-mmlab/mmsegmentation.git
cd mmsegmentation
pip install -v -e .
# '-v' 表示详细模式，更多的输出
# '-e' 表示以可编辑模式安装工程，
# 因此对代码所做的任何修改都生效，无需重新安装

然后就安装成功了，我们用pycharm打开项目文件，就可以进行对应的操作了。

安装过程中如果有不解的地方，可以参考官方文档：开始：安装和运行 MMSeg — MMSegmentation 1.2.2 文档

3.mmsegmentation使用（代码结构介绍）

mmsegmentation下载下来后的主要代码结构如下：

# MMSegmentation 算法库目录结构的主要部分 
mmsegmentation 
   | 
   |- configs                        # 配置文件 
   |     |- _base_                   ## 配置文件 
   |     |     |- datasets             ### 数据集相关配置文件 
   |     |     |- models               ### 模型相关配置文件 
   |     |     |- schedules            ### 不同迭代次数下的训练配置文件 
   |     |     |- default_runtime.py   ### 运行相关的默认的设置 
   |     |- ann                      ## 各个模型对应的配置加载文件
   |     |- apcnet    
   |     |- ...                         
   |- data                           # 数据集存放文件夹 
   |- mmseg  
   |     |- models                    ## 分割模型具体实现代码 
   |     |     |- backbones             ### 主干网络 
   |     |     |- decode_heads          ### 解码头 
   |     |     |- losses                ### 损失函数 
   |     |     |- necks                 ### 颈 
   |     |     |- segmentors            ### 构建完整分割网络的代码 
   |     |     |- utils                 ### 构建模型时的辅助工具 
   |- tools  
   |     |- train.py                  ## 训练脚本 
   |     |- test.py                   ## 测试脚本 
   |     |- ...                       
   |- ...

可以得知的是，mmsegmentation下所有关于模型、数据集、训练次数等东西的配置都在configs文件夹下。也就是说，如果我们只是使用现成的模型，那我们只需要关注configs文件夹就好。

比如models文件夹点进去会有对应的模型配置文件，里面对模型的训练方式、骨干网络、语义分割类别数等信息进行了配置。

4.mmsegmentation使用实战（deeplab v3+为例）

4.1配置

下面对使用mmsegmentation中的deeplab v3+模型，在voc2012数据集中进行语义分割：

首先，我们下载voc2012数据集，并将其放到data文件夹下：

然后我们创建myModel/deeplabv3+文件夹，用来存放模型配置文件。

然后我们到configs文件夹中找到对应的配置文件，需要的配置文件主要有5个：

（1）configs/_base_/datasets/下的pascal_voc12.py文件，里面会对数据集中的存储文件进行了配置，如果数据集按照上述文件路径存放的不需要更改，直接复制到myModel/deeplabv3+文件夹中。否则，需要修改配置文件中的对应路径：

（2）configs/_base_/models/下的deeplabv3plus_r50-d8.py文件，里面对deeplabv3+模型的具体信息进行了定义。如果使用GPU训练及预测的就不用修改，使用CPU训练及预测的则改为BN：

（3）configs/_base_/下的schedule_20k.py文件，该文件定义了训练迭代次数。

（4）configs/_base_/下的default_runtime.py文件，该文件中需要对模型的预训练权重文件进行加载：

选择None则表示不加载预训练文件。

注意：每一个模型的预训练权重文件都可以在configs文件夹下的对应模型名称文件夹的readme文件中找到下载地址。如deeplabv3+的预训练文件在configs/deeplabv3plus/下的readme文件中：

找到对应的数据集、骨干网络、迭代次数下的[model]对应的地址即可下载。

下载完成后，我们可以创建对应的weights文件夹存放对应的权重文件。（参考上图中default_runtime.py的文件路径）

（5）configs/deeplabv3plus/下的deeplabv3plus_r50-d8_4xb4-20k_voc12aug-512x512.py文件。

文件名称格式为模型名称_骨干网络_迭代次数_数据集，对应上述各文件的选择。因此根据文件名称，我们需要确定好对应的数据集、模型、及迭代文件。

该文件对如何加载各类配置文件的路径进行了配置，是最终的集成配置文件。我们需要对里面的文件路径进行配置：

（6）然后

最终，我们再创建一个work-dir文件夹，存放训练时的文件。

最终文件夹下目录格式如下：

4.2训练：

配置完成后，我们就可以进行训练了。

找到tools文件夹下的train.py文件，修改里面对应文件的路径：

然后运行，既可以进行训练。

然后，我们可以看到对应的checkpoint文件被保存在哪里：

在work-dir文件夹中会有对应的日期文件夹，里面存放对应的数据。

4.3预测：

我们可以使用训练后得到的权重文件，或者直接使用预训练权重文件进行预测：

import numpy as np
import matplotlib.pyplot as plt

from mmseg.apis import init_model, inference_model, show_result_pyplot
import mmcv
import cv2

# 模型 config 配置文件
config_file = 'E:/tc-project/mmsegmentation/myModel/deeplabv3+/deeplabv3plus_r50-d8_4xb4-20k_voc12aug-512x512.py'

# 模型 checkpoint 权重文件
checkpoint_file = 'E:/tc-project/mmsegmentation/myModel/deeplabv3+/weights/deeplabv3plus_r50-d8_512x512_40k_voc12aug.pth'

device = 'cpu'
model = init_model(config_file, checkpoint_file, device=device)

# 测试图片路径
img_path = 'C:/Users/pc/Desktop/test.jpg'
img_bgr = cv2.imread(img_path)
result = inference_model(model, img_bgr)
pred_mask = result.pred_sem_seg.data[0].cpu().numpy()
cmap = plt.cm.colors.ListedColormap(['black', 'pink'])

# 显示语义分割结果
plt.figure(figsize=(14, 8))
plt.subplot(1,2,1)
plt.imshow(img_bgr[:,:,::-1])
plt.axis('off')
plt.subplot(1,2,2)
plt.imshow(img_bgr[:,:,::-1])
plt.imshow(pred_mask, cmap=cmap, alpha=0.6) # alpha 高亮区域透明度，越小越接近原图
plt.axis('off')
plt.savefig('outputs/deeplab_1.jpg')
plt.show()

上述代码中的配置文件是指最终的集成配置文件，权重文件本次使用了预训练权重文件。

运行效果：