LRASPP_MobileNetv3训练自己的分割数据集

treelly

于 2024-12-26 17:30:03 发布

阅读量916

点赞数 27

分类专栏： python 文章标签：计算机视觉

本文链接：https://blog.csdn.net/treelly/article/details/144745546

版权

python 专栏收录该内容

7 篇文章

订阅专栏

今年做的布料瑕疵检测项目，其中一个需求是准确识别布料边缘。之前只针对白色布料，使用opencv进行边缘检测，也够用了，随着布料颜色越来越多样，采集图片的光源环境越来越多变，opencv边缘检测模块已经不能满足需求。同事建议我试试语义分割模型的效果，并推荐了一个轻量级模型：LRASPP_MobileNetV3。
这篇文章用来记录使用LRASPP_MobileNetV3训练自己的分割数据集的实战过程。从环境搭建、数据准备、模型训练、再到导出trt模型和测试，方便日后查阅。

0 基础知识

论文地址：Searching for MobileNetV3
Pytorch 实现代码：
https://github.com/WZMIAOMIAO/deep-learning-for-image-processing/tree/master/pytorch_segmentation/lraspp
MobileNet v3 LR-ASPP

0.1 语义分割和实例分割

计算机视觉任务中，我容易混淆的是：目标检测、语义分割、实例分割、全景分割（任务结果依次对应下图，图片来自这里图解AI：图像分割）
目标检测：识别图像中存在的内容和检测其位置，如下图，以识别和检测人（person）为例
语义分割：对图像中的每个像素打上类别标签，如下图，把图像分为人（红色）、树木（深绿）、草地（浅绿）、天空（蓝色）标签
实例分割：目标检测+语义分割。对比右上图、左下图，如以人（person）为目标，语义分割不区分属于相同类别的不同实例（所有人都标为红色），实例分割区分同类的不同实例（使用不同颜色区分不同的人）
全景分割：语义分割+实例分割，即要对所有目标都检测出来，又要区分出同个类别中的不同实例。
在这里插入图片描述

0.2 了解LRASPP_MobileNetv3

MobileNet

要使用一个模型，首先要大致了解这个模型是什么，主要干什么用，再到怎么用。刚看到这一串英文名称时，我是一脸懵的，即使阅读了一些MobileNetv3相关的文章，仍然云里雾里，直到看到这篇文章（从MobileNet看轻量级神经网络的发展），有种拨开云雾的顿悟。
引用一下人家文章中的图哈
计算机视觉领域内的卷积神经网络模型发展历程：LeNet——>AlexNet——>VGG——>ResNet，深度学习网络模型在图像处理中应用的效果越来越好。神经网络体积越来越大，结构越来越复杂，预测和训练需要的硬件资源也逐步增多，往往只能在高算力的服务器中运行。
为了适应移动设备和嵌入式设备的算力限制，深度学习领域开始发展轻量级网络模型。2016年至今，出现了SqueezeNet、ShuffleNet、NasNet、MnasNet和MobileNet等模型，它们在尽量保持准确率的同时，体积更小、速度更快。其中，MobileNet尤为突出，谷歌在2019年推出了性能更好的MobileNetV3。
看到下面的对比数据，不得不感叹确实牛。

使用Google Pixel-1手机测试，最新版MobileNetV3-Large运行时间达到66ms，GoogleNet运行速度约为250ms，而VGG-16由于一次性需要加载至内存的空间已超过500MB，手机系统会报内存溢出错误导致无法运行。

MobileNet可以在移动终端实现的应用，包括检测，分割，分类，人脸识别等。

LRASPP

LRASPP：Lite Reduced Atrous Spatial Pyramid Pooling，轻量化简减少空洞空间金字塔池化
LRASPP是 MobileNetV3 中用于语义分割任务的一种高效分割解码器。

1、环境搭建

win10+Anaconda+PyCharm+ISAT(一个好用的半自动分割标注工具，安装及使用参考实例分割标注工具ISAT安装及使用教程)
1、下载项目代码（选择一个，同事发给我的代码是他更改过的）
https://github.com/WZMIAOMIAO/deep-learning-for-image-processing/tree/master/pytorch_segmentation/lraspp
MobileNet v3 LR-ASPP
2、准备环境
创建虚拟环境

conda create -n lraspp python=3.8
conda activate lraspp

安装pytorch

通过pip安装的pytorch默认是cpu版本，gpu版本的pytorch需去 pytorch 官网手动安装。

切换到项目文件夹下，安装需要的库：

cd LRASPP_MobileNetV3
pip install -r requirements.txt

2、制作数据集

2.1 图片标注

使用ISAT对images文件夹中的图片进行标注，生成jsons，数据转换成yolo格式（生成masks图片）
对于模型训练来说，有用的是images、masks文件夹。
在这里插入图片描述

2.2 划分训练集、验证集

划分训练集、验证集后的文件夹结构如下，images中是原始图片，labels中是masks图片
在这里插入图片描述示例代码：

def split_data(src_dir, image_dir, label_dir, split_ratio=0.9):
    """
    数据集划分为训练集、验证集
    :param src_dir: 原始数据集文件夹，包含图像和标签
    :param image_dir: 划分后的图片文件夹，包含train\val子文件夹
    :param txt_dir: 划分后的标签文件夹，包含train\val子文件夹
    :param split_ratio: 训练集、验证集划分比例，默认为0.9
    :return: None
    """
    # 创建数据文件夹目录
    new_dir_list = {'img_train': os.path.join(image_dir, 'train'),
                    'img_val': os.path.join(image_dir, 'val'),
                    'label_train': os.path.join(label_dir, 'train'),
                    'label_val': os.path.join(label_dir, 'val')}
    os.makedirs(new_dir_list['img_train'], exist_ok=True)
    os.makedirs(new_dir_list['img_val'], exist_ok=True)
    os.makedirs(new_dir_list['label_train'], exist_ok=True)
    os.makedirs(new_dir_list['label_val'], exist_ok=True)

    origin_img_dir = os.path.join(src_dir, 'images')
    origin_label_dir = os.path.join(src_dir, 'masks')

    total_img = os.listdir(origin_img_dir)
    total_label = os.listdir(origin_label_dir)

    total_num = len(total_img)
    tv = int(total_num * split_ratio)
    list_index = range(total_num)
    train_index = random.sample(list_index, tv)

    for i in list_index:
        if i in train_index:
            new_img_dir = new_dir_list['img_train']
            new_label_dir = new_dir_list['label_train']
        else:
            new_img_dir = new_dir_list['img_val']
            new_label_dir = new_dir_list['label_val']

        img_file = os.path.join(origin_img_dir, total_img[i])
        label_file = os.path.join(origin_label_dir, total_label[i])
        new_img_file = os.path.join(new_img_dir, total_img[i])
        new_label_file = os.path.join(new_label_dir, total_label[i])
        # os.rename(img_file, new_img_file)  #移动
        # os.rename(label_file, new_label_file)
        shutil.copy(img_file, new_img_file)  # 复制
        shutil.copy(label_file, new_label_file)
    print(f"数据划分完成，训练集{tv},测试集{total_num-tv}")
if __name__ == "__main__":
    src_dir = r'.\01resized'
    image_dir = r'.\02trainModel\images'
    label_dir = r'.\02trainModel\labels'
    split_data(src_dir, image_dir, label_dir)

3、模型训练

在开始训练之前，代码中修改这些内容：分割类的数量（背景默认是一类）、类名称，图片大小【后续填坑】
我的训练命令：

python train.py --batch 10 --height 100 --width 800 --lr 0.05 --epochs 100 --out boundary --data ./02trainModel
# 我的训练图片大小是100*800, --out: 生成模型保存路径 --data: 训练数据路径

模型训练效果：
在这里插入图片描述
450张训练图片，从上图可以看出，验证集结果比训练集结果低0.2左右，明显是模型过拟合了；验证集指标突然下降又上升，是什么原因呢，清楚的小伙伴请留言赐教。
过拟合可能的原因：

训练数据量过小：460张图片，这对于深度学习模型来说是相对较小的。
模型复杂度：如果模型的参数量过大，可能会导致其在训练集上表现良好，但无法在验证集上泛化。模型可能学到了一些特定于训练数据的噪声，而不是泛化到未见过的数据。
学习率过大或过小：如果学习率选择不当，模型可能会在训练过程中出现不稳定的训练结果，或在验证集上表现较差。
缺乏正则化：如果没有采取正则化措施（如 Dropout、L2 正则化等），模型容易过拟合训练数据。
数据增强（Data Augmentation）缺乏：如果没有对训练数据进行适当的增强，模型可能无法有效泛化到验证集。

4、导出模型

1、首先运行export_onnx脚本，从.pt模型生成onnx模型
2、从onnx模型导出trt模型

trtexec.exe --onnx=segment.onnx --saveEngine=segment.trt --buildOnly --workspace=1024 --fp16