Giter VIP home page Giter VIP logo

pytorch-yolov4's Introduction

Pytorch-YOLOv4

A minimal PyTorch implementation of YOLOv4.

├── README.md
├── dataset.py            dataset
├── demo.py               demo to run pytorch --> tool/darknet2pytorch
├── demo_darknet2onnx.py  tool to convert into onnx --> tool/darknet2pytorch
├── demo_pytorch2onnx.py  tool to convert into onnx
├── models.py             model for pytorch
├── train.py              train models.py
├── cfg.py                cfg.py for train
├── cfg                   cfg --> darknet2pytorch
├── data            
├── weight                --> darknet2pytorch
├── tool
│   ├── camera.py           a demo camera
│   ├── coco_annotatin.py       coco dataset generator
│   ├── config.py
│   ├── darknet2pytorch.py
│   ├── region_loss.py
│   ├── utils.py
│   └── yolo_layer.py

image

0. Weights Download

0.1 darkent

0.2 pytorch

you can use darknet2pytorch to convert it yourself, or download my converted model.

1. Train

use yolov4 to train your own data

  1. Download weight

  2. Transform data

    For coco dataset,you can use tool/coco_annotatin.py.

    # train.txt
    image_path1 x1,y1,x2,y2,id x1,y1,x2,y2,id x1,y1,x2,y2,id ...
    image_path2 x1,y1,x2,y2,id x1,y1,x2,y2,id x1,y1,x2,y2,id ...
    ...
    ...
    
  3. Train

    you can set parameters in cfg.py.

     python train.py -g [GPU_ID] -dir [Dataset direction] ...
    

2. Inference

  • Load the pretrained darknet model and darknet weights to do the inference
python demo.py -cfgfile <cfgFile> -weightfile <weightFile> -imgfile <imgFile>
  • Load pytorch weights (pth file) to do the inference
python models.py <num_classes> <weightfile> <imgfile> <namefile(optional)>

3. Darknet2ONNX (Evolving)

  • This script is to convert the official pretrained darknet model into ONNX

  • Pytorch version Recommended: 1.4.0

  • Install onnxruntime

    pip install onnxruntime
  • Run python script to generate onnx model and run the demo

    python demo_darknet2onnx.py <cfgFile> <weightFile> <imageFile> <batchSize>

    This script will generate 2 onnx models.

    • One is for running the demo (batch_size=1)
    • The other one is what you want to generate (batch_size=batchSize)

4. Pytorch2ONNX (Evolving)

  • You can convert your trained pytorch model into ONNX using this script

  • Pytorch version Recommended: 1.4.0

  • Install onnxruntime

    pip install onnxruntime
  • Run python script to generate onnx model and run the demo

    python demo_pytorch2onnx.py <weight_file> <image_path> <batch_size> <n_classes> <IN_IMAGE_H> <IN_IMAGE_W>

    For example:

    python demo_pytorch2onnx.py yolov4.pth dog.jpg 8 80 416 416

    This script will generate 2 onnx models.

    • One is for running the demo (batch_size=1)
    • The other one is what you want to generate (batch_size=batch_size)

5. ONNX2TensorRT (Evolving)

  • TensorRT version Recommended: 7.0, 7.1

  • Run the following command to convert VOLOv4 onnx model into TensorRT engine

    trtexec --onnx=<onnx_file> --explicitBatch --saveEngine=<tensorRT_engine_file> --workspace=<size_in_megabytes> --fp16
    • Note: If you want to use int8 mode in conversion, extra int8 calibration is needed.
  • Run the demo (this demo here only works when batchSize=1)

    python demo_trt.py <tensorRT_engine_file> <input_image> <input_H> <input_W>
    • Note1: input_H and input_W should agree with the input size in the original darknet cfg file as well as the latter onnx file.
    • Note2: extra NMS operations are needed for the tensorRT output. This demo uses TianXiaomo's NMS code from tool/utils.py.

6. ONNX2Tensorflow

Reference:

@article{yolov4,
  title={YOLOv4: YOLOv4: Optimal Speed and Accuracy of Object Detection},
  author={Alexey Bochkovskiy, Chien-Yao Wang, Hong-Yuan Mark Liao},
  journal = {arXiv},
  year={2020}
}

pytorch-yolov4's People

Contributors

cmdbug avatar dhiyu avatar ersheng-ai avatar tehtea avatar tianxiaomo avatar vcasecnikovs avatar

Watchers

 avatar

Recommend Projects

  • React photo React

    A declarative, efficient, and flexible JavaScript library for building user interfaces.

  • Vue.js photo Vue.js

    🖖 Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.

  • Typescript photo Typescript

    TypeScript is a superset of JavaScript that compiles to clean JavaScript output.

  • TensorFlow photo TensorFlow

    An Open Source Machine Learning Framework for Everyone

  • Django photo Django

    The Web framework for perfectionists with deadlines.

  • D3 photo D3

    Bring data to life with SVG, Canvas and HTML. 📊📈🎉

Recommend Topics

  • javascript

    JavaScript (JS) is a lightweight interpreted programming language with first-class functions.

  • web

    Some thing interesting about web. New door for the world.

  • server

    A server is a program made to process requests and deliver data to clients.

  • Machine learning

    Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.

  • Game

    Some thing interesting about game, make everyone happy.

Recommend Org

  • Facebook photo Facebook

    We are working to build community through open source technology. NB: members must have two-factor auth.

  • Microsoft photo Microsoft

    Open source projects and samples from Microsoft.

  • Google photo Google

    Google ❤️ Open Source for everyone.

  • D3 photo D3

    Data-Driven Documents codes.