yolov5-tensorrt

Add video inference .

Code fork form https://github.com/wang-xinyu/tensorrtx.git .

The Pytorch implementation is ultralytics/yolov5.

How to Run

1. generate yolov5s.wts from pytorch implementation with yolov5s.pt

git clone https://github.com/wang-xinyu/tensorrtx.git
git clone https://github.com/ultralytics/yolov5.git
// download its weights 'yolov5s.pt'
cd yolov5
cp ../tensorrtx/yolov5s/gen_wts.py .
python gen_wts.py
// a file 'yolov5s.wts' will be generated.

2. put yolov5s.wts into yolov5, build and run

mv yolov5s.wts ../tensorrtx/yolov5/
cd ../tensorrtx/yolov5
mkdir build
cd build
cmake ..
make
sudo ./yolov5s -s             // serialize model to plan file i.e. 'yolov5s.engine'
sudo ./yolov5s -d  ../samples // deserialize plan file and run inference, the images in samples will be processed.

3. check the images generated, as follows. _zidane.jpg and _bus.jpg

Config

• Input shape defined in yololayer.h
• Number of classes defined in yololayer.h
• FP16/FP32 can be selected by the macro in yolov5s.cpp
• GPU id can be selected by the macro in yolov5s.cpp
• NMS thresh in yolov5s.cpp
• BBox confidence thresh in yolov5s.cpp
• Batch size in yolov5s.cpp

More Information

See the readme in home page.