Are you sure you want to delete this task? Once this task is deleted, it cannot be recovered.
Eric Sheng 4ccef0ec8f | 3 years ago | |
---|---|---|
DeepStream | 3 years ago | |
cfg | 3 years ago | |
data | 3 years ago | |
tool | 3 years ago | |
.gitignore | 3 years ago | |
License.txt | 3 years ago | |
README.md | 3 years ago | |
Use_yolov4_to_train_your_own_data.md | 3 years ago | |
cfg.py | 3 years ago | |
dataset.py | 3 years ago | |
demo.py | 3 years ago | |
demo_darknet2onnx.py | 3 years ago | |
demo_pytorch2onnx.py | 3 years ago | |
demo_tensorflow.py | 3 years ago | |
demo_trt.py | 3 years ago | |
evaluate_on_coco.py | 3 years ago | |
models.py | 3 years ago | |
requirements.txt | 3 years ago | |
train.py | 3 years ago |
A minimal PyTorch implementation of YOLOv4.
Paper Yolo v4: https://arxiv.org/abs/2004.10934
Source code:https://github.com/AlexeyAB/darknet
More details: http://pjreddie.com/darknet/yolo/
├── 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_annotation.py coco dataset generator
│ ├── config.py
│ ├── darknet2pytorch.py
│ ├── region_loss.py
│ ├── utils.py
│ └── yolo_layer.py
you can use darknet2pytorch to convert it yourself, or download my converted model.
use yolov4 to train your own data
Download weight
Transform data
For coco dataset,you can use tool/coco_annotation.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 ...
...
...
Train
you can set parameters in cfg.py.
python train.py -g [GPU_ID] -dir [Dataset direction] ...
ONNX and TensorRT models are converted from Pytorch (TianXiaomo): Pytorch->ONNX->TensorRT.
See following sections for more details of conversions.
Model type | AP | AP50 | AP75 | APS | APM | APL |
---|---|---|---|---|---|---|
DarkNet (YOLOv4 paper) | 0.471 | 0.710 | 0.510 | 0.278 | 0.525 | 0.636 |
Pytorch (TianXiaomo) | 0.466 | 0.704 | 0.505 | 0.267 | 0.524 | 0.629 |
TensorRT FP32 + BatchedNMSPlugin | 0.472 | 0.708 | 0.511 | 0.273 | 0.530 | 0.637 |
TensorRT FP16 + BatchedNMSPlugin | 0.472 | 0.708 | 0.511 | 0.273 | 0.530 | 0.636 |
Model type | AP | AP50 | AP75 | APS | APM | APL |
---|---|---|---|---|---|---|
DarkNet (YOLOv4 paper) | 0.412 | 0.628 | 0.443 | 0.204 | 0.444 | 0.560 |
Pytorch (TianXiaomo) | 0.404 | 0.615 | 0.436 | 0.196 | 0.438 | 0.552 |
TensorRT FP32 + BatchedNMSPlugin | 0.412 | 0.625 | 0.445 | 0.200 | 0.446 | 0.564 |
TensorRT FP16 + BatchedNMSPlugin | 0.412 | 0.625 | 0.445 | 0.200 | 0.446 | 0.563 |
Image input size is NOT restricted in 320 * 320
, 416 * 416
, 512 * 512
and 608 * 608
.
You can adjust your input sizes for a different input ratio, for example: 320 * 608
.
Larger input size could help detect smaller targets, but may be slower and GPU memory exhausting.
height = 320 + 96 * n, n in {0, 1, 2, 3, ...}
width = 320 + 96 * m, m in {0, 1, 2, 3, ...}
Load the pretrained darknet model and darknet weights to do the inference (image size is configured in cfg file already)
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> <IN_IMAGE_H> <IN_IMAGE_W> <namefile(optional)>
Load converted ONNX file to do inference (See section 3 and 4)
Load converted TensorRT engine file to do inference (See section 5)
There are 2 inference outputs.
[batch, num_boxes, 1, 4]
which represents x1, y1, x2, y2 of each bounding box.[batch, num_boxes, num_classes]
indicating scores of all classes for each bounding box.Until now, still a small piece of post-processing including NMS is required. We are trying to minimize time and complexity of post-processing.
This script is to convert the official pretrained darknet model into ONNX
Pytorch version Recommended:
Install onnxruntime
pip install onnxruntime
Run python script to generate ONNX model and run the demo
python demo_darknet2onnx.py <cfgFile> <weightFile> <imageFile> <batchSize>
You can convert your trained pytorch model into ONNX using this script
Pytorch version Recommended:
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
Run the following command to convert YOLOv4 ONNX model into TensorRT engine
trtexec --onnx=<onnx_file> --explicitBatch --saveEngine=<tensorRT_engine_file> --workspace=<size_in_megabytes> --fp16
Run the following command to convert YOLOv4 ONNX model into TensorRT engine
trtexec --onnx=<onnx_file> \
--minShapes=input:<shape_of_min_batch> --optShapes=input:<shape_of_opt_batch> --maxShapes=input:<shape_of_max_batch> \
--workspace=<size_in_megabytes> --saveEngine=<engine_file> --fp16
For example:
trtexec --onnx=yolov4_-1_3_320_512_dynamic.onnx \
--minShapes=input:1x3x320x512 --optShapes=input:4x3x320x512 --maxShapes=input:8x3x320x512 \
--workspace=2048 --saveEngine=yolov4_-1_3_320_512_dynamic.engine --fp16
python demo_trt.py <tensorRT_engine_file> <input_image> <input_H> <input_W>
This demo here only works when batchSize is dynamic (1 should be within dynamic range) or batchSize=1, but you can update this demo a little for other dynamic or static batch sizes.
Note1: input_H and input_W should agree with the input size in the original ONNX file.
Note2: extra NMS operations are needed for the tensorRT output. This demo uses python NMS code from tool/utils.py
.
First:Conversion to ONNX
tensorflow >=2.0
1: Thanks:github:https://github.com/onnx/onnx-tensorflow
2: Run git clone https://github.com/onnx/onnx-tensorflow.git && cd onnx-tensorflow
Run pip install -e .
Note:Errors will occur when using "pip install onnx-tf", at least for me,it is recommended to use source code installation
cd DeepStream
make
For single batch,
trtexec --onnx=<onnx_file> --explicitBatch --saveEngine=<tensorRT_engine_file> --workspace=<size_in_megabytes> --fp16
For multi-batch,
trtexec --onnx=<onnx_file> --explicitBatch --shapes=input:Xx3xHxW --optShapes=input:Xx3xHxW --maxShapes=input:Xx3xHxW --minShape=input:1x3xHxW --saveEngine=<tensorRT_engine_file> --fp16
Note :The maxShapes could not be larger than model original shape.
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}
}
No Description
Python C++ INI Text Cuda other
Dear OpenI User
Thank you for your continuous support to the Openl Qizhi Community AI Collaboration Platform. In order to protect your usage rights and ensure network security, we updated the Openl Qizhi Community AI Collaboration Platform Usage Agreement in January 2024. The updated agreement specifies that users are prohibited from using intranet penetration tools. After you click "Agree and continue", you can continue to use our services. Thank you for your cooperation and understanding.
For more agreement content, please refer to the《Openl Qizhi Community AI Collaboration Platform Usage Agreement》