mmsegmentation 训练自己的数据集

一. MMSegmentation是什么？

MMSegmentation 是一个基于 PyTorch 的语义分割开源工具箱，它是 OpenMMLab 项目的一部分。他与MMDetection类似，集成了各种语义分割算法，可以快速验证语义分割效果。

二. 环境准备

参考：https://mmsegmentation.readthedocs.io/en/latest/get_started.html#verification

1. 安装环境

# 创建虚拟环境并进入
conda create --name openmmlab python=3.8 -y
conda activate openmmlab

# 安装pytorch时，要根据自己的cuda版本进行安装，比如我的cuda版本是10.2
conda install pytorch==1.12.0 torchvision==0.13.0 torchaudio==0.12.0 cudatoolkit=10.2 -c pytorch

# 用MIM安装MMCV
pip install -U openmim
mim install mmengine
mim install "mmcv>=2.0.0"

# 安装MMSegmentation
git clone -b main https://github.com/open-mmlab/mmsegmentation.git
cd mmsegmentation
pip install -v -e .

2. 验证环境

# 下载config和checkpoint文件
mim download mmsegmentation --config pspnet_r50-d8_4xb2-40k_cityscapes-512x1024 --dest .

# 验证推理效果
python demo/image_demo.py demo/demo.png configs/pspnet/pspnet_r50-d8_4xb2-40k_cityscapes-512x1024.py pspnet_r50-d8_512x1024_40k_cityscapes_20200605_003338-2966598c.pth --device cuda:0 --out-file result.jpg

三. 数据准备

建议直接将自己的数据制作成VOC的数据格式，跟主流数据格式一致。
VOC数据下载地址：https://github.com/StoneWST/Dataset-Tool-Segmentation/

简单说下，JPEGImages是原始图片，SegmentationClass存储的是.png格式的标注数据，读取label数据的时候，对于单通道数据，像素值要和类别的id 相等。比如我的数’background’,‘building’在单通道的label数据中要像素值分别要为background’= 0，‘building’= 1。

四. 修改配置文件

我习惯用 python /tools/train 某个标准config.py 指令先生成一个标准的config文件，然后复制出来再改，有点麻烦，但是后边省事。而且我个人建议初学者重新生成一个config文件，而不是自己直接改标注的。

python tools/train.py configs/deeplabv3/deeplabv3-r50-d8512x51220k_voc12aug.py

运行这个指令肯定会出错，但是在work_dirs下会生成一个对应的文件夹，然后在这个文件夹下会有一个自动生成的config文件。
config文件大概长这样：

很多很繁杂，但是不要害怕，我们需要改动的地方并不多。记住：一数据，二类别，三训练配置。
注意，我记得原始VOC还进行了数据增强，你把带有Aug对应的地方以此删掉即可。
这些地方都要改成你的数据集对应的地址。后面的val，test道理亦然。

修改mmseg/datasets/voc.py文件，找到voc.py（因为我们的格式是VOC），做如下修改。其中PALETTE你可以简单理解为颜色，background对应[0, 0, 0]就是黑色，依次类推。

./mmseg/utils/class_names.py 做类似修改：

注意：每次修改mmseg内部的代码后要重新编译。编译命令如下：

pip install -v -e .

开始训练模型：

CUDA_VISIBLE_DEVICES=1,2,3,4 python tools/train.py deeplabv3plus_r18-d8_4xb4-20k_myproject-512x512.py

按照教程来，你的工程就顺利训练起来了, 如下图：

五. 测试

由于我的任务是二分类，测试代码由demo/image_demo.py修改过来。

from argparse import ArgumentParser

from mmengine.model import revert_sync_batchnorm


if __name__ == '__main__':
    import os
    import glob
    import cv2
    import numpy as np
    config = "your_config_path.py"
    checkpoint = "your_checkpoint_path.pth"
    device = 'cuda:0'
    model = init_model(config, checkpoint, device)
    img_list = glob.glob("./images/*.jpg")
    for img in img_list:
        result = inference_model(model, img)
        result = result.pred_sem_seg.data                    # 取出分割结果
        result = np.array(result.cpu()[0,:,:])*255
        # single to 3 channels
        result = np.expand_dims(result, axis=2)
        result = np.concatenate((result, result, result), axis=-1)
        print("inference success...")
        image = cv2.imread(img)
        output_img = np.hstack([image, result])
        out_file = os.path.join("outputs", os.path.basename(img))
        cv2.imwrite(out_file, output_img)

到此，就介绍完基于MMSegmentatoin训练自己的数据集，有问题欢迎留言讨论。