LibtorchSegmentation

英文 | 中文

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

该库的主要特点是：

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

📚 Libtorch 教程 📚

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

📋 目录

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

⏳ 快速开始

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")

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

💡 示例

• 使用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"/*要显示的类名*/);

预测结果如下所示：

🧑‍🚀 训练自己的数据

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

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

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

📦 模型

架构

• Unet [论文]
• FPN [论文]
• PAN [论文]
• PSPNet [论文]
• LinkNet [论文]
• DeepLabV3 [论文]
• DeepLabV3+ [论文]

编码器

• ResNet
• ResNext
• VGG

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

🛠 安装

依赖项：

• Opencv 3+
• Libtorch 1.7+

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

⏳ 待办事项

• 更多分割架构和骨干网络
  • UNet++ [论文]
  • ResNest
  • Se-Net
  • ...
• 数据增强
  • 随机水平翻转
  • 随机垂直翻转
  • 随机尺度旋转
  • ...
• 训练技巧
  • 结合Dice和交叉熵损失
  • 冻结骨干网络
  • 多步学习率调度
  • ...

🤝 致谢

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

• 官方pytorch
• qubvel SMP
• wkentaro labelme
• nlohmann json

📝 引用

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

🛡️ 许可证

项目以MIT许可证分发。

相关仓库

基于libtorch，我发布了以下仓库：

• LibtorchTutorials
• LibtorchSegmentation
• LibtorchDetection

最后但同样重要的是，别忘了给个star...

欢迎提交问题或拉取请求，我们需要贡献者。