dhiyu / pytorch-YOLOv4

Minimal PyTorch implementation of YOLOv4

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Pytorch-YOLOv4

A minimal PyTorch implementation of YOLOv4.

├── README.md
├── dataset.py       dataset
├── demo.py          demo to run pytorch --> tool/darknet2pytorch
├── darknet2onnx.py  tool to convert into onnx --> tool/darknet2pytorch
├── demo_onnx.py     demo to run the converted onnx model
├── 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.Weight

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

python demo.py <cfgFile> <weightFile> <imgFile>

3.Darknet2ONNX

  • Install onnxruntime

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

    python demo_onnx.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.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}
}

About

Minimal PyTorch implementation of YOLOv4

License:Apache License 2.0


Languages

Language:Python 100.0%