从零开始复现seaformer(语义分割)训练自己的数据集——linux

引言：

出于模型轻量化需求，需对原有的皮肤分割模型进行重新研发。seaformer是作为今年复旦大学和腾讯联合提出的轻量级语义分割算法，具有很好的参考价值。因此，作者基于seaformer算法对自研数据集进行训练，完成轻量级皮肤分割模型的开发。

一、seaformer源码下载

1.下载地址：https://github.com/fudan-zvg/SeaFormer

2.解压到自己的目录当中，我们主要使用sesaformer-seg，具体结构如下：

二、制作相关数据集

1.使用labelme工具对自己的数据进行标注，具体不再赘述

2.通过labelme工具进行数据标注后，我们可以得到原图和.json文件

3.将.josn文件批量转换为用于训练的mask，文件具体格式如下：

批量转换可参考：将labelme的json文件批量转为png形式 - 知乎 (zhihu.com)https://zhuanlan.zhihu.com/p/649269470创建data文件夹，在data文件夹下新建cityscapes文件夹，cityscapes文件夹底下再新建gtFine与leftImg8bit文件夹（gtFine与leftImg8bit同级）。在leftImg8bit文件夹中存放原图(后缀为.png格式)，gtFine文件夹用来存放8bit的全黑mask图像（此时暂为空）。具体结构如下：

.json文件完成转换后，第一张全黑的图片就是我们需要的图像，参照seaformer中训练数据集cityscapes，要求的标签图像格式为8bit（原标签图像为24bit），通过以下代码实现转换:

import cv2
import os

# 24bit to 8bit
bit24_dir = r'E:\Face_Parsing_BiseNetV2-master\ATR\masks'  # 存放.json转换后的24bit全黑图像文件夹
if not os.path.exists('data/cityscapes/gtFine'):
    os.makedirs('data/cityscapes/gtFine')
bit8_dir = 'data/cityscapes/gtFine'
png_names = os.listdir(bit24_dir)
for i in png_names:
    img = cv2.imread(bit24_dir + '/' + i)
    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    cv2.imencode('.png', gray)[1].tofile(bit8_dir + '/' + i)

4.创建并运行train_val.py完成训练数据集与测试数据集的划分

import os
import random
import shutil

total_list = []
train_list = []
val_list = []

image_path = 'data/cityscapes/leftImg8bit'
label_path = 'data/cityscapes/gtFine'

# 清空
for dir in ['train', 'val']:
    image_dir = os.path.join(image_path, dir)
    label_dir = os.path.join(label_path, dir)
    if os.path.exists(image_dir):
        shutil.rmtree(image_dir)
    os.makedirs(image_dir)
    if os.path.exists(label_dir):
        shutil.rmtree(label_dir)
    os.makedirs(label_dir)

for root, dirs, files in os.walk(image_path):
    for file in files:
        if file.endswith('png'):
            total_list.append(file)

total_size = len(total_list)
train_size = int(total_size * 0.8)
val_size = total_size - train_size

train_list = random.sample(total_list, train_size)
remain_list = list(set(total_list) - set(train_list))
val_list = random.sample(remain_list, val_size)

for file in total_list:
    image_path_0 = os.path.join(image_path, file)
    label_file = file.split('.')[0] + '.png'
    label_path_0 = os.path.join(label_path, label_file)
    if file in train_list:
        image_path_1 = os.path.join(image_path, 'train', file)
        shutil.move(image_path_0, image_path_1)

        label_path_1 = os.path.join(label_path, 'train', label_file)
        shutil.move(label_path_0, label_path_1)

    elif file in val_list:
        image_path_1 = os.path.join(image_path, 'val', file)
        shutil.move(image_path_0, image_path_1)

        label_path_1 = os.path.join(label_path, 'val', label_file)
        shutil.move(label_path_0, label_path_1)

print(len(total_list))
print(len(train_list))
print(len(val_list))

生成目录格式：

三、训练seaformer网络

1.环境搭建：

1）conda 虚拟环境创建

conda create -n your_env_name python=x.x  #python==3.8(作者版本)

2）激活虚拟环境并安装pytorch（gpu），建议安装1.8版本，太高版本的pytorch在安装mmcv-full时会在build时出现报错

3）安装mmcv-full

pip install -U openmim
mim install mmcv-full==1.5.0 #seaformer要求1.3.14版本（1.5.0为作者版本）

4）安装相关依赖包

cd /home/uboot/chen/SeaFormer-main/seaformer-seg
pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple/

2.修改相关代码文件

1）修改SeaFormer-main\seaformer-seg\local_configs\seaformer\seaformer_small_1024x512_160k_1x8city.py文件

将num_classes修改为自己的类别数，background算一个类，其他按照自己的图像尺寸进行修改。

2）在seaformer-seg下创建modelzoos文件夹，在modelzoos文件夹下创建classification文件夹用于存放seaformer提供的预训练权重文件，将下载的预训练权重改名为SeaFormer_S.pth。修改SeaFormer-main\seaformer-seg\local_configs\seaformer\seaformer_small.py，改为自己的类别数并将第一行中的SyncBN改为BN

3）修改SeaFormer-main\seaformer-seg\mmseg\datasets\cityscapes.py,修改为自己的标签名，颜色随便填;img_suffix和seg_map_suffix记得修改不然会出现加载不到图像报错的情况。

四、开始训练

bash tools/dist_train.sh local_configs/seaformer/seaformer_small_1024x512_160k_1x8city.py 1

迭代次数可以在SeaFormer-main\seaformer-seg\local_configs\_base_\schedules\schedule_160k.py中设置

五、结果输出

记录一下，若存在不对的地方，请大家多多包涵~