PyTorch ,ONNX and TensorRT implementation of YOLOv4

Overview

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_annotation.py       coco dataset generator
│   ├── config.py
│   ├── darknet2pytorch.py
│   ├── region_loss.py
│   ├── utils.py
│   └── yolo_layer.py

image

0. Weights Download

0.1 darknet

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_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 ...
    ...
    ...
    
  3. Train

    you can set parameters in cfg.py.

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

2. Inference

2.1 Performance on MS COCO dataset (using pretrained DarknetWeights from https://github.com/AlexeyAB/darknet)

ONNX and TensorRT models are converted from Pytorch (TianXiaomo): Pytorch->ONNX->TensorRT. See following sections for more details of conversions.

  • val2017 dataset (input size: 416x416)
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
  • testdev2017 dataset (input size: 416x416)
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

2.2 Image input size for inference

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, ...}

2.3 Different inference options

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

2.4 Inference output

There are 2 inference outputs.

  • One is locations of bounding boxes, its shape is [batch, num_boxes, 1, 4] which represents x1, y1, x2, y2 of each bounding box.
  • The other one is scores of bounding boxes which is of shape [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.

3. Darknet2ONNX

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

  • Pytorch version Recommended:

    • Pytorch 1.4.0 for TensorRT 7.0 and higher
    • Pytorch 1.5.0 and 1.6.0 for TensorRT 7.1.2 and higher
  • Install onnxruntime

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

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

3.1 Dynamic or static batch size

  • Positive batch size will generate ONNX model of static batch size, otherwise, batch size will be dynamic
    • Dynamic batch size will generate only one ONNX model
    • Static batch size will generate 2 ONNX models, one is for running the demo (batch_size=1)

4. Pytorch2ONNX

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

  • Pytorch version Recommended:

    • Pytorch 1.4.0 for TensorRT 7.0 and higher
    • Pytorch 1.5.0 and 1.6.0 for TensorRT 7.1.2 and higher
  • 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

4.1 Dynamic or static batch size

  • Positive batch size will generate ONNX model of static batch size, otherwise, batch size will be dynamic
    • Dynamic batch size will generate only one ONNX model
    • Static batch size will generate 2 ONNX models, one is for running the demo (batch_size=1)

5. ONNX2TensorRT

  • TensorRT version Recommended: 7.0, 7.1

5.1 Convert from ONNX of static Batch size

  • 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
    • Note: If you want to use int8 mode in conversion, extra int8 calibration is needed.

5.2 Convert from ONNX of dynamic Batch size

  • 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

5.3 Run the demo

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.

6. ONNX2Tensorflow

7. ONNX2TensorRT and DeepStream Inference

  1. Compile the DeepStream Nvinfer Plugin
    cd DeepStream
    make 
  1. Build a TRT Engine.

For single batch,

trtexec --onnx= --explicitBatch --saveEngine= --workspace= --fp16

For multi-batch,

trtexec --onnx= --explicitBatch --shapes=input:Xx3xHxW --optShapes=input:Xx3xHxW --maxShapes=input:Xx3xHxW --minShape=input:1x3xHxW --saveEngine= --fp16

Note :The maxShapes could not be larger than model original shape.

  1. Write the deepstream config file for the TRT Engine.

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}
}
Owner
DL CV OCR and algorithm optimization
The object detection pipeline is based on Ultralytics YOLOv5

AYOLOv2 The main goal of this repository is to rewrite the object detection pipeline with a better code structure for better portability and adaptabil

153 Dec 22, 2022
Implementation of SETR model, Original paper: Rethinking Semantic Segmentation from a Sequence-to-Sequence Perspective with Transformers.

SETR - Pytorch Since the original paper (Rethinking Semantic Segmentation from a Sequence-to-Sequence Perspective with Transformers.) has no official

zhaohu xing 112 Dec 16, 2022
U-Net for GBM

My Final Year Project(FYP) In National University of Singapore(NUS) You need Pytorch(stable 1.9.1) Both cuda version and cpu version are OK File Str

PinkR1ver 1 Oct 27, 2021
Files for a tutorial to train SegNet for road scenes using the CamVid dataset

SegNet and Bayesian SegNet Tutorial This repository contains all the files for you to complete the 'Getting Started with SegNet' and the 'Bayesian Seg

Alex Kendall 800 Dec 31, 2022
Codebase for ECCV18 "The Sound of Pixels"

Sound-of-Pixels Codebase for ECCV18 "The Sound of Pixels". *This repository is under construction, but the core parts are already there. Environment T

Hang Zhao 318 Dec 20, 2022
HandTailor: Towards High-Precision Monocular 3D Hand Recovery

HandTailor This repository is the implementation code and model of the paper "HandTailor: Towards High-Precision Monocular 3D Hand Recovery" (arXiv) G

Lv Jun 113 Jan 06, 2023
[NeurIPS 2020] This project provides a strong single-stage baseline for Long-Tailed Classification, Detection, and Instance Segmentation (LVIS).

A Strong Single-Stage Baseline for Long-Tailed Problems This project provides a strong single-stage baseline for Long-Tailed Classification (under Ima

Kaihua Tang 514 Dec 23, 2022
NEATEST: Evolving Neural Networks Through Augmenting Topologies with Evolution Strategy Training

NEATEST: Evolving Neural Networks Through Augmenting Topologies with Evolution Strategy Training

Göktuğ Karakaşlı 16 Dec 05, 2022
Code for the paper "Adversarial Generator-Encoder Networks"

This repository contains code for the paper "Adversarial Generator-Encoder Networks" (AAAI'18) by Dmitry Ulyanov, Andrea Vedaldi, Victor Lempitsky. Pr

Dmitry Ulyanov 279 Jun 26, 2022
RARA: Zero-shot Sim2Real Visual Navigation with Following Foreground Cues

RARA: Zero-shot Sim2Real Visual Navigation with Following Foreground Cues FGBG (foreground-background) pytorch package for defining and training model

Klaas Kelchtermans 1 Jun 02, 2022
PrimitiveNet: Primitive Instance Segmentation with Local Primitive Embedding under Adversarial Metric (ICCV 2021)

PrimitiveNet Source code for the paper: Jingwei Huang, Yanfeng Zhang, Mingwei Sun. [PrimitiveNet: Primitive Instance Segmentation with Local Primitive

Jingwei Huang 47 Dec 06, 2022
This repository provides data for the VAW dataset as described in the CVPR 2021 paper titled "Learning to Predict Visual Attributes in the Wild"

Visual Attributes in the Wild (VAW) This repository provides data for the VAW dataset as described in the CVPR 2021 Paper: Learning to Predict Visual

Adobe Research 36 Dec 30, 2022
SMPLpix: Neural Avatars from 3D Human Models

subject0_validation_poses.mp4 Left: SMPL-X human mesh registered with SMPLify-X, middle: SMPLpix render, right: ground truth video. SMPLpix: Neural Av

Sergey Prokudin 292 Dec 30, 2022
"Reinforcement Learning for Bandit Neural Machine Translation with Simulated Human Feedback"

This is code repo for our EMNLP 2017 paper "Reinforcement Learning for Bandit Neural Machine Translation with Simulated Human Feedback", which implements the A2C algorithm on top of a neural encoder-

Khanh Nguyen 131 Oct 21, 2022
[CVPR 2021 Oral] ForgeryNet: A Versatile Benchmark for Comprehensive Forgery Analysis

ForgeryNet: A Versatile Benchmark for Comprehensive Forgery Analysis ForgeryNet: A Versatile Benchmark for Comprehensive Forgery Analysis [arxiv|pdf|v

Yinan He 78 Dec 22, 2022
A PyTorch implementation of the baseline method in Panoptic Narrative Grounding (ICCV 2021 Oral)

A PyTorch implementation of the baseline method in Panoptic Narrative Grounding (ICCV 2021 Oral)

Biomedical Computer Vision @ Uniandes 52 Dec 19, 2022
Pip-package for trajectory benchmarking from "Be your own Benchmark: No-Reference Trajectory Metric on Registered Point Clouds", ECMR'21

Map Metrics for Trajectory Quality Map metrics toolkit provides a set of metrics to quantitatively evaluate trajectory quality via estimating consiste

Mobile Robotics Lab. at Skoltech 31 Oct 28, 2022
Towards Flexible Blind JPEG Artifacts Removal (FBCNN, ICCV 2021)

Towards Flexible Blind JPEG Artifacts Removal (FBCNN, ICCV 2021)

Jiaxi Jiang 282 Jan 02, 2023
Python module providing a framework to trace individual edges in an image using Gaussian process regression.

Edge Tracing using Gaussian Process Regression Repository storing python module which implements a framework to trace individual edges in an image usi

Jamie Burke 7 Dec 27, 2022