PyTorch图像分割模型——segmentation_models_pytor

一、概要
是一个基于的图像分割神经网络
这个新集合由俄罗斯的程序员小哥Pavel 一手打造,对于图像分割而言简直就是神器般的存在 。
地址:
该库的主要功能有:
高级API(只需两行即可创建神经网络);用于二分类和多类分割的9种模型架构(包括传奇的Unet)(Unet、Unet++、MAnet、、FPN、、PAN、、+);每种架构有113种可用的编码器;所有编码器均具有预训练的权重,以实现更快更好的收敛 。
这篇文章我主要介绍UNet++模型的训练和检测,UNet在图像分割中已经足够传奇,可以利用底层的特征(同分辨率级联)改善上采样的信息不足,对训练所要求的数据集并不是特别大,而且准确率也相当的高,UNet++又在UNet的基础上增加了剪枝的操作配合深监督,整合不同层次的特征以提升精度,且减少了图像分割所用的时间,只是在内存的要求上相应的增加了些 。关于UNet和UNet++具体的解读可参考这两篇文章:
UNET详解和UNET++介绍(零基础)
[论文解读]UNet++解读 + 它是如何对UNet改进 + 作者的研究态度和方式
话不多说,我们开始基于库进行UNet++模型训练 。
二、环境配置
首先我们需创建一个虚拟环境,我这里采用的是.7版本 。
【PyTorch图像分割模型——segmentation_models_pytor】conda create -n segmentation python==3.7
创建好后进入虚拟环境,并安装我们这里要用的图像分割神器-- 。
pip install segmentation-models-pytorch
安装--会一同安装上torch和,但是这时要注意了,这里安装进去的是CPU版的而且是最新版的,如果你确实打算用cpu来做的话那后面的也可以忽略,但是做目标检测和图像分割不用GPU简直没有灵魂啊,而且也为了方便我们同目标检测的神经网络结合,所以我们在--安装好后要先卸载掉torch和,然后再安装相应的cuda版本的torch和 。我电脑采用的是.0,所以首先下载后对应.0的torch和,并进行安装 。我这里采用的是下载好的离线安装包 。
pip install torch-1.7.1+cu110-cp37-cp37m-win_amd64.whlpip install torchvision-0.8.2+cu110-cp37-cp37m-win_amd64.whl
安装好后虚拟环境中的如下图所示 。
同时,我们还需安装以下库:
pip install albumentationspip install matplotlibpip install imageiopip install opencv-python
安装好后,虚拟环境中所含的库如下图所示 。
此时,我们的前期工作就准备好了,接下来开始图像分割模型的训练和测试 。
三、库之UNet++模型训练 3.1 UNet++训练
首先我们看一下库是怎么使用他所包含的9种模型的 。
由于该库是基于框架构建的,因此创建的细分模型只是一个 nn.,可以轻松地创建它 。
import segmentation_models_pytorch as smpmodel = smp.UnetPlusPlus(encoder_name="resnet34",# 选择解码器, 例如 mobilenet_v2 或 efficientnet-b7encoder_weights="imagenet",# 使用预先训练的权重imagenet进行解码器初始化in_channels=1,# 模型输入通道(1个用于灰度图像,3个用于RGB等)classes=3,# 模型输出通道(数据集所分的类别总数))
如果要使用UNet、FPN等模型,只要改成model = smp.Unet()或model = smp.FPN()即可 。
所有模型均具有预训练的编码器,因此必须按照权重预训练的相同方法准备数据 。
from segmentation_models_pytorch.encoders import get_preprocessing_fnpreprocess_input = get_preprocessing_fn('resnet18', pretrained='imagenet')