Skip to content

C++ trainable detection library based on libtorch (or pytorch c++). Yolov4 tiny provided now.

License

Notifications You must be signed in to change notification settings

AllentDan/LibtorchDetection

Repository files navigation

logo C++ Library with Neural Networks for Object
Detection Based on LibTorch.

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

📋 Table of content

  1. Examples
  2. Train your own data
  3. Installation
  4. To do list
  5. Thanks
  6. Citing
  7. License
  8. Related repository

💡 Examples

  1. Download the VOC image dataset, and split the dataset into train and val parts as follows:
dataset
├── train
│   ├── images
|   |     ├──xxx.jpg
|   |     └......
│   ├── labels
|   |     ├──xxx.xml
|   |     └......
├── val
│   ├── images
|   |     ├──zzz.jpg
|   |     └......
│   ├── labels
|   |     ├──zzz.xml
|   |     └......
  1. Download the yolo4_tiny pretrained weight here. And load it into your cpp project as follows:
Detector detector;
detector.Initialize(-1, /*gpu id, -1 for cpu*/
                    416, /*resize width*/
                    416, /*resize height*/
                    "your path to class name.txt");
detector.Train("your path to dataset dir", 
                    ".jpg", /*image type*/
                    30,/*training epochs*/
                    4, /*batch size*/
                    0.001, /*learning rate*/
                    "path to save detector.pt",
                    "path to load pretrained yolo4_tiny.pt");
  1. Predicting test. A detector.pt file is provided in the project here (trained on VOC for one epoch, just for testing...). Click and download, then you can directly test the detection result through:
cv::Mat image = cv::imread("your path to 2007_005331.jpg");
Detector detector;
detector.Initialize(0, 416, 416, "your path to voc_classes.txt");
detector.LoadWeight("detector.pt"/*the saved .pt path*/);
detector.Predict(image, true,/*show result or not*/, 0.1, /*confidence thresh*/, 0.3/*nms thresh*/);

the predicted result shows as follow:

🧑‍🚀 Train your own data

  • Create your own dataset. Using labelImg through "pip install" and label your images. Split the output xml files and images into folders just like the example above.

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

🛠 Installation

Dependency:

Windows:

Configure the environment for libtorch development. Visual studio and Qt Creator are verified for libtorch1.7+.

Linux && MacOS:

Install libtorch and opencv.

For libtorch, follow the official pytorch c++ tutorials here.

For opencv, follow the official opencv install steps here.

If you have already configured them both, congratulations!!! Download the pretrained weight here and a demo .pt file here into weights. Then just do the following:

export Torch_DIR='/path/to/libtorch'
cd build
cmake ..
make
./LibtorchDetection

⏳ ToDo

  • More detection architectures, mainly one-stage algorithms.
  • Data augmentations.
  • Training tricks.

🤝 Thanks

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

📝 Citing

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

🛡️ License

Project is distributed under MIT License.

Related repository

Based on libtorch, I released following repositories:

Last but not least, don't forget your star...

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