LibtorchSegmentation

英文 | 中文

logo
基于LibTorch的图像分割神经网络C++库。

</div>

⭐如果这个项目对您有帮助，请给个星星⭐

该库的主要特点是：

高级API（只需一行代码即可创建神经网络）
7种用于二分类和多分类分割的模型架构（包括传奇的Unet）
15种可用编码器
所有编码器都有预训练权重，可实现更快更好的收敛
与pytorch cuda相比，推理速度提升35%或更多，CPU速度相同。（在RTX 2070s上测试Unet）

📚 Libtorch 教程 📚

如果您想了解更多关于Libtorch Segment库的信息，请访问Libtorch教程项目。

📋 目录

快速开始
示例
训练自己的数据
模型
1. 架构
2. 编码器
安装
致谢
待办事项
引用
许可证
相关仓库

⏳ 快速开始 <a name="start"></a>

1. 使用Libtorch Segment创建您的第一个分割模型

这里提供了一个resnet34的torchscript文件下载链接。分割模型只是一个LibTorch torch::nn::Module，创建起来非常简单：

#include "Segmentor.h"
auto model = UNet(1, /*类别数*/
                  "resnet34", /*编码器名称，可以是resnet50或其他*/
                  "resnet34.pt的路径"/*在ImageNet上预训练的权重路径，由torchscript生成*/
                  );

查看表格了解可用的模型架构
查看表格了解可用的编码器及其对应的权重

2. 生成您自己的预训练权重

所有编码器都有预训练权重。以与预训练权重相同的方式准备数据可能会得到更好的结果（更高的指标分数和更快的收敛）。您还可以只训练解码器和分割头，同时冻结主干网络。

import torch
from torchvision import models

# 以resnet34为例
model = models.resnet34(pretrained=True)
model.eval()
var=torch.ones((1,3,224,224))
traced_script_module = torch.jit.trace(model, var)
traced_script_module.save("resnet34.pt")

恭喜！您已完成！现在您可以使用您喜欢的主干网络和分割框架来训练您的模型了。

💡 示例 <a name="examples"></a>

使用PASCAL VOC数据集的图像训练人物分割模型。"voc_person_seg"目录包含32个用于训练的json标签及其对应的jpeg图像，以及8个用于验证的json标签及其对应图像。

Segmentor<FPN> segmentor;
segmentor.Initialize(0/*gpu id，-1表示使用cpu*/,
                    512/*调整宽度*/,
                    512/*调整高度*/,
                    {"background","person"}/*类名字典，包括背景*/,
                    "resnet34"/*主干网络名称*/,
                    "resnet34.pt的路径");
segmentor.Train(0.0003/*初始学习率*/,
                300/*训练轮数*/,
                4/*批次大小*/,
                "voc_person_seg的路径",
                ".jpg"/*图像类型*/,
                "保存segmentor.pt的路径");

预测测试。项目中提供了一个segmentor.pt文件下载链接。它是通过带有ResNet34主干的FPN训练了几个轮次得到的。您可以直接通过以下方式测试分割结果：

cv::Mat image = cv::imread("voc_person_seg\\val\\2007_004000.jpg的路径");
Segmentor<FPN> segmentor;
segmentor.Initialize(0,512,512,{"background","person"},
                      "resnet34","resnet34.pt的路径");
segmentor.LoadWeight("segmentor.pt"/*保存的.pt文件路径*/);
segmentor.Predict(image,"person"/*要显示的类名*/);

预测结果如下所示：

🧑‍🚀 训练自己的数据 <a name="trainingOwn"></a>

创建自己的数据集。通过"pip install"安装labelme并标注你的图像。将输出的json文件和图像分别放入文件夹，如下所示：

数据集
├── 训练
│   ├── xxx.json
│   ├── xxx.jpg
│   └......
├── 验证
│   ├── xxxx.json
│   ├── xxxx.jpg
│   └......

训练或测试。与"voc_person_seg"示例类似，将"voc_person_seg"替换为你自己的数据集路径。
参考训练技巧来提高最终的训练效果。

📦 模型 <a name="models"></a>

架构 <a name="architectures"></a>

编码器 <a name="encoders"></a>

ResNet
ResNext
VGG

以下是Libtorch Segment支持的编码器列表。除resnest外，所有编码器权重都可以通过torchvision生成。选择适当的编码器系列并点击展开表格，选择特定的编码器及其预训练权重。

<details> <summary style="margin-left: 25px;">ResNet</summary> <div style="margin-left: 25px;">

编码器	权重	参数量, M
resnet18	imagenet	11M
resnet34	imagenet	21M
resnet50	imagenet	23M
resnet101	imagenet	42M
resnet152	imagenet	58M

</div> </details> <details> <summary style="margin-left: 25px;">ResNeXt</summary> <div style="margin-left: 25px;">

编码器	权重	参数量, M
resnext50_32x4d	imagenet	22M
resnext101_32x8d	imagenet	86M

</div> </details> <details> <summary style="margin-left: 25px;">ResNeSt</summary> <div style="margin-left: 25px;">

编码器	权重	参数量, M
timm-resnest14d	imagenet	8M
timm-resnest26d	imagenet	15M
timm-resnest50d	imagenet	25M
timm-resnest101e	imagenet	46M
timm-resnest200e	imagenet	68M
timm-resnest269e	imagenet	108M
timm-resnest50d_4s2x40d	imagenet	28M
timm-resnest50d_1s4x24d	imagenet	23M

</div> </details> <details> <summary style="margin-left: 25px;">SE-Net</summary> <div style="margin-left: 25px;">

编码器	权重	参数量, M
senet154	imagenet	113M
se_resnet50	imagenet	26M
se_resnet101	imagenet	47M
se_resnet152	imagenet	64M
se_resnext50_32x4d	imagenet	25M
se_resnext101_32x4d	imagenet	46M

</div> </details> <details> <summary style="margin-left: 25px;">VGG</summary> <div style="margin-left: 25px;">

编码器	权重	参数量, M
vgg11	imagenet	9M
vgg11_bn	imagenet	9M
vgg13	imagenet	9M
vgg13_bn	imagenet	9M
vgg16	imagenet	14M
vgg16_bn	imagenet	14M
vgg19	imagenet	20M
vgg19_bn	imagenet	20M

</div> </details>

🛠 安装 <a name="installation"></a>

依赖项：

Windows：

配置libtorch开发环境。Visual studio和Qt Creator已验证可用于libtorch1.7x版本。

Linux && MacOS：

安装libtorch和opencv。

对于libtorch，请参照官方pytorch c++教程此处。

对于opencv，请按照官方opencv安装步骤此处进行操作。

如果您已经配置好了这两个，恭喜您！请在此处下载预训练权重，并在此处下载演示用.pt文件到weights文件夹。

构建共享或静态库 -DBUILD_SHARED=<TRUE/FALSE>:

export Torch_DIR='/path/to/libtorch'
cd build
cmake -DBUILD_SHARED=TRUE ..
make
sudo make install

构建测试：

cd test
mkdir build && cd build
cmake ..
make
./resnet34 ../../voc_person_seg/val/2007_003747.jpg ../../weights/resnet34.pt ../../weights/segmentor.pt

⏳ 待办事项 <a name="todo"></a>

🤝 致谢 <a name="thanks"></a>

目前，这些项目提供了很大帮助。

📝 引用

@misc{Chunyu:2021,
  Author = {Chunyu Dong},
  Title = {Libtorch Segment},
  Year = {2021},
  Publisher = {GitHub},
  Journal = {GitHub repository},
  Howpublished = {\url{https://github.com/AllentDan/SegmentationCpp}}
}