[English](https://github.com/AllentDan/SegmentationCpp) | 中文 <div align="center"> ![logo](https://raw.githubusercontent.com/AllentDan/ImageBase/main/OpenSource/LibtorchSegment.png) **基于[LibTorch](https://pytorch.org/)的C++开源图像分割神经网络库.** </div> **⭐如果有用请给我一个star⭐** 这个库具有以下优点: - 高级的API (只需一行代码就可创建网络) - 7 种模型架构可用于单类或者多类的分割任务 (包括Unet) - 15 种编码器网络 - 所有的编码器都有预训练权重,可以更快更好地收敛 - 相比于python下的GPU前向推理速度具有30%或以上的提速, cpu下保持速度一致. (Unet测试于RTX 2070S). ### [📚 Libtorch教程 📚](https://github.com/AllentDan/LibtorchTutorials) 如果你想对该开源项目有更多更详细的了解,请前往本人另一个开源项目:[Libtorch教程](https://github.com/AllentDan/LibtorchTutorials) . ### 📋 目录 1. [快速开始](#start) 2. [例子](#examples) 3. [训练自己的数据](#trainingOwn) 4. [模型](#models) 1. [架构](#architectures) 2. [编码器](#encoders) 5. [安装](#installation) 6. [ToDo](#todo) 7. [感谢](#thanks) 8. [引用](#citing) 9. [证书](#license) 10. [相关项目](#related_repos) ### ⏳ 快速开始 <a name="start"></a> #### 1. 用 Libtorch Segment 创建你的第一个分割网络 [这](https://github.com/AllentDan/LibtorchSegmentation/releases/download/weights/segmentor.pt)是一个resnet34的torchscript模型,可以作为骨干网络权重。分割模型是 LibTorch 的 torch::nn::Module的派生类, 可以很容易生成: ```cpp #include "Segmentor.h" auto model = UNet(1, /*num of classes*/ "resnet34", /*encoder name, could be resnet50 or others*/ "path to resnet34.pt"/*weight path pretrained on ImageNet, it is produced by torchscript*/ ); ``` - 见 [表](#architectures) 查看所有支持的模型架构 - 见 [表](#encoders) 查看所有的编码器网络和相应的预训练权重 #### 2. 生成自己的预训练权重 所有编码器均具有预训练的权重。加载预训练权重,以相同的方式训练数据,可能会获得更好的结果(更高的指标得分和更快的收敛速度)。还可以在冻结主干的同时仅训练解码器和分割头。 ```python import torch from torchvision import models # resnet50 for example model = models.resnet50(pretrained=True) model.eval() var=torch.ones((1,3,224,224)) traced_script_module = torch.jit.trace(model, var) traced_script_module.save("resnet50.pt") ``` 恭喜你! 大功告成! 现在,您可以使用自己喜欢的主干和分割框架来训练模型了。 ### 💡 例子 <a name="examples"></a> - 使用来自PASCAL VOC数据集的图像进行人体分割数据训练模型. "voc_person_seg" 目录包含32个json标签及其相应的jpeg图像用于训练,还有8个json标签以及相应的图像用于验证。 ```cpp Segmentor<FPN> segmentor; segmentor.Initialize(0/*gpu id, -1 for cpu*/, 512/*resize width*/, 512/*resize height*/, {"background","person"}/*class name dict, background included*/, "resnet34"/*backbone name*/, "your path to resnet34.pt"); segmentor.Train(0.0003/*initial leaning rate*/, 300/*training epochs*/, 4/*batch size*/, "your path to voc_person_seg", ".jpg"/*image type*/, "your path to save segmentor.pt"); ``` - 预测测试。项目中提供了以ResNet34为骨干网络的FPN网络,训练了一些周期得到segmentor.pt文件[在这](https://github.com/AllentDan/LibtorchSegmentation/releases/download/weights/segmentor.pt)。 您可以直接测试分割结果: ```cpp cv::Mat image = cv::imread("your path to voc_person_seg\\val\\2007_004000.jpg"); Segmentor<FPN> segmentor; segmentor.Initialize(0,512,512,{"background","person"}, "resnet34","your path to resnet34.pt"); segmentor.LoadWeight("segmentor.pt"/*the saved .pt path*/); segmentor.Predict(image,"person"/*class name for showing*/); ``` 预测结果显示如下: ![](https://raw.githubusercontent.com/AllentDan/SegmentationCpp/main/prediction.jpg) ### 🧑🚀 训练自己的数据 <a name="trainingOwn"></a> - 创建自己的数据集. 使用"pip install"安装[labelme](https://github.com/wkentaro/labelme)并标注你的图像. 将输出的json文件和图像分成以下文件夹: ``` Dataset ├── train │ ├── xxx.json │ ├── xxx.jpg │ └...... ├── val │ ├── xxxx.json │ ├── xxxx.jpg │ └...... ``` - 训练或测试。就像“ voc_person_seg”的示例一样,用自己的数据集路径替换“ voc_person_seg”。 - 记得使用[训练技巧](https://github.com/AllentDan/LibtorchSegmentation/blob/main/docs/training%20tricks.md)以提高模型的训练效果。 ### 📦 Models <a name="models"></a> #### Architectures <a name="architectures"></a> - [x] Unet [[paper](https://arxiv.org/abs/1505.04597)] - [x] FPN [[paper](http://presentations.cocodataset.org/COCO17-Stuff-FAIR.pdf)] - [x] PAN [[paper](https://arxiv.org/abs/1805.10180)] - [x] PSPNet [[paper](https://arxiv.org/abs/1612.01105)] - [x] LinkNet [[paper](https://arxiv.org/abs/1707.03718)] - [x] DeepLabV3 [[paper](https://arxiv.org/abs/1706.05587)] - [x] DeepLabV3+ [[paper](https://arxiv.org/abs/1802.02611)] #### Encoders <a name="encoders"></a> - [x] ResNet - [x] ResNext - [x] VGG 以下是该项目中受支持的编码器的列表。除resnest外,所有编码器权重都可以通过torchvision生成。选择适当的编码器,然后单击以展开表格,然后选择特定的编码器及其预训练的权重。 <details> <summary style="margin-left: 25px;">ResNet</summary> <div style="margin-left: 25px;"> |Encoder |Weights |Params, 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;"> |Encoder |Weights |Params, M | |--------------------------------|:------------------------------:|:------------------------------:| |resnext50_32x4d |imagenet |22M | |resnext101_32x8d |imagenet |86M | </div> </details> <details> <summary style="margin-left: 25px;">ResNeSt</summary> <div style="margin-left: 25px;"> |Encoder |Weights |Params, 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;"> |Encoder |Weights |Params, 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;"> |Encoder |Weights |Params, 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> **依赖库:** - [Opencv 3+](https://opencv.org/releases/) - [Libtorch 1.7+](https://pytorch.org/) **Windows:** 配置libtorch 开发环境. [Visual studio](https://allentdan.github.io/2020/12/16/pytorch%E9%83%A8%E7%BD%B2torchscript%E7%AF%87) 和 [Qt Creator](https://allentdan.github.io/2021/01/21/QT%20Creator%20+%20Opencv4.x%20+%20Libtorch1.7%E9%85%8D%E7%BD%AE/#more)已经通过libtorch1.7x release的验证. **Linux && MacOS:** 安装libtorch和opencv。 对于libtorch, 按照官方[教程](https://pytorch.org/tutorials/advanced/cpp_export.html)安装。 对于opencv, 按照官方安装[步骤](https://github.com/opencv/opencv)。 如果你都配置好了他们,恭喜!!! 下载一个resnet34的预训练权重,[点击下载](https://github.com/AllentDan/LibtorchSegmentation/releases/download/weights/resnet34.pt)和一个示例.pt文件,[点击下载](https://github.com/AllentDan/LibtorchSegmentation/releases/download/weights/segmentor.pt),放入weights文件夹。 更改src/main.cpp中的图片路径预训练权重和加载的segmentor权重路径。随后,build路径在终端输入: ```bash export Torch_DIR='/path/to/libtorch' cd build cmake .. make ./LibtorchSegmentation ``` ### ⏳ ToDo <a name="todo"></a> - [ ] 更多的骨干网络和分割框架 - [ ] UNet++ [[paper](https://arxiv.org/pdf/1807.10165.pdf)] - [ ] ResNest - [ ] Se-Net - [ ] ... - [x] 数据增强 - [x] 随机水平翻转 - [x] 随机垂直翻转 - [x] 随机缩放和旋转 - [ ] ... - [x] 训练技巧 - [x] 联合损失:dice和交叉熵 - [x] 冻结骨干网络 - [x] 学习率衰减策略 - [ ] ... ### 🤝 感谢 <a name="thanks"></a> 以下是目前给予帮助的项目. - [official pytorch](https://github.com/pytorch/pytorch) - [qubvel SMP](https://github.com/qubvel/segmentation_models.pytorch) - [wkentaro labelme](https://github.com/wkentaro/labelme) - [nlohmann json](https://github.com/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}} } ``` ### 🛡️ 证书 <a name="license"></a> 该项目以 [MIT License](https://github.com/qubvel/segmentation_models.pytorch/blob/master/LICENSE)开源, ## 相关项目 <a name="related_repos"></a> 基于libtorch,我释放了如下开源项目: - [LibtorchTutorials](https://github.com/AllentDan/LibtorchTutorials) - [LibtorchSegmentation](https://github.com/AllentDan/LibtorchSegmentation) - [LibtorchDetection](https://github.com/AllentDan/LibtorchDetection) 别忘了点赞哟 ![stargazers over time](https://starchart.cc/AllentDan/SegmentationCpp.svg)