English | 中文

If the star number equal to 2n-1, please make it 2n.

C++ library with Neural Networks for Image
Segmentation based on LibTorch.

The main features of this library are:

• High level API (just a line to create a neural network)
• 7 models architectures for binary and multi class segmentation (including legendary Unet)
• 7 available encoders
• All encoders have pre-trained weights for faster and better convergence
• 2x or more faster than pytorch cuda inferece, same speed for cpu. (Unet tested in gtx 2070s).

📚 Libtorch Tutorials 📚

Visit Libtorch Tutorials Project if you want to know more about Libtorch Segment library.

📋 Table of content

1. Quick start
2. Examples
3. Train your own data
4. Models
   1. Architectures
   2. Encoders
5. Installation
6. Thanks
7. Citing
8. License

⏳ Quick start

1. Create your first Segmentation model with Libtorch Segment

Segmentation model is just a LibTorch torch::nn::Module, which can be created as easy as:

#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*/
                  );

• see table with available model architectures
• see table with available encoders and their corresponding weights

2. Generate your own pretrained weights

All encoders have pretrained weights. Preparing your data the same way as during weights pre-training may give your better results (higher metric score and faster convergence). And you can also train only the decoder and segmentation head while freeze the backbone.

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

Congratulations! You are done! Now you can train your model with your favorite backbone and segmentation framework.

💡 Examples

• Training model for person segmentation using images from PASCAL VOC Dataset. "voc_person_seg" dir contains 32 json labels and their corresponding jpeg images for training and 8 json labels with corresponding images for validation.

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

• Predicting test. A segmentor.pt file is provided in the project. It is trained through a FPN with ResNet34 backbone for a few epochs. You can directly test the segmentation result through:

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*/);

the predicted result shows as follow:

🧑‍🚀 Train your own data

• Create your own dataset. Using labelme through "pip install" and label your images. Split the output json files and images into folders just like below:

Dataset
├── train
│   ├── xxx.json
│   ├── xxx.jpg
│   └......
├── val
│   ├── xxxx.json
│   ├── xxxx.jpg
│   └......

• Training or testing. Just like the example of "voc_person_seg", replace "voc_person_seg" with your own dataset path.

📦 Models

Architectures

• Unet [paper]
• FPN [paper]
• PAN [paper]
• PSPNet [paper]
• LinkNet [paper]
• DeepLabV3 [paper]
• DeepLabV3+ [paper]
• UNet++ [paper]

Encoders

• ResNet
• ResNext
• ResNest
• Se-Net

The following is a list of supported encoders in the Libtorch Segment. All the encoders weights can be generated through torchvision except resnest. Select the appropriate family of encoders and click to expand the table and select a specific encoder and its pre-trained weights.

ResNet

Encoder	Weights	Params, M
resnet18	imagenet	11M
resnet34	imagenet	21M
resnet50	imagenet	23M
resnet101	imagenet	42M
resnet152	imagenet	58M

ResNeXt

Encoder	Weights	Params, M
resnext50_32x4d	imagenet	22M
resnext101_32x8d	imagenet	86M

ResNeSt

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

SE-Net

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

🛠 Installation

Dependency:

• Opencv 3+
• Libtorch 1.7+

Windows:

Configure the environment for libtorch development. Visual studio and Qt Creator are verified for libtorch1.7x release. Only Visual Studio configuration blogs provided english version by now, Qt english version ASAP.

Linux && MacOS:

Follow the official pytorch c++ tutorials here. It can be no more difficult than windows.

🤝 Thanks

This project is under developing. By now, these projects helps a lot.

• official pytorch
• qubvel SMP
• wkentaro labelme
• nlohmann json

📝 Citing

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

🛡️ License

Project is distributed under MIT License. Last but not least, don't forget your star...

Feel free to commit issues or pull requests, contributors wanted.