Skip to content

fish-kong/Yolov5-obb-Tensorrt-Infer

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

4 Commits
 
 
 
 
 
 

Repository files navigation

yolov5_for_oriented_object_detection

The Pytorch implementation is yolov5_obb

actually ,This repo is based on yolov5 and only used at Linux PC (partial dependency install hardly at windows PC)

How to Run, yolov5s as example

1.train your datasets in {yolov5_obb} generate .pt at {yolov5_obb/runs/exp[?]/weights/best.pt}

  1. generate .wts from {yolov5_obb/runs/exp[?]/weights/best.pt}
cp gen_wts.py {yolov5_obb}
python gen_wts.py -w runs/train/exp[]/weights/best.pt -o yolov5s.wts
// a file 'yolov5s.wts' will be generated.
  1. build **tensorrt engine **
cd {yolov5_obb_tensorrt_cpp}
// update CLASS_NUM/Inputsize in yololayer.h 
// caution:CLASS_NUM= your classes +180(180 for angle classes)
mkdir build
cd build
cp {yolov5_obb}/yolov5s.wts ../yolov5s.wts
cmake ..
make
sudo ./yolov5_gen -s [.wts] [.engine] [n/s/m/l/x]  // serialize model to plan file
// For example yolov5s
sudo ./yolov5_gen -s ../yolov5s.wts ../yolov5s.engine s
  1. use tensorr engine
sudo ./yolov5_use ../yolov5s.engine ../images/test.jpg

image