fix yolop tensorrt, add usage; 修复 yolop tensorrt 编译问题, 添加使用说明

toolkits/deploy/CMakeLists.txt

@@ -8,12 +8,18 @@ option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
 set(CMAKE_CXX_STANDARD 11)
 set(CMAKE_BUILD_TYPE Release)
 
-
-find_package(ZED 3 REQUIRED)
-find_package(CUDA ${ZED_CUDA_VERSION} EXACT REQUIRED)
+find_package(ZED 3)
+if(ZED_FOUND)
+    find_package(CUDA ${ZED_CUDA_VERSION} EXACT REQUIRED)
+else(ZED_FOUND)
+    find_package(CUDA  REQUIRED)
+endif(ZED_FOUND)
 
 include_directories(${PROJECT_SOURCE_DIR}/include)
 
+find_package(OpenCV REQUIRED)
+include_directories(${OpenCV_INCLUDE_DIRS})
+
 # cuda
 include_directories(/usr/local/cuda-10.2/include)
 link_directories(/usr/local/cuda-10.2/lib64)
@@ -26,20 +32,24 @@ link_directories(/usr/local/zed/lib)
 
 set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -Ofast -Wfatal-errors -D_MWAITXINTRIN_H_INCLUDED")
 
-set(ZED_LIBS ${ZED_LIBRARIES} ${CUDA_CUDA_LIBRARY} ${CUDA_CUDART_LIBRARY})
-
-coda_add_library(myplugins SHARED ${PROJECT_SOURCE_DIR}/yololayer.cu)
+# to generate plugins
+cuda_add_library(myplugins SHARED ${PROJECT_SOURCE_DIR}/yololayer.cu)
 target_link_libraries(myplugins nvinfer cudart)
 
-find_package(OpenCV REQUIRED)
-include_directories(${OpenCV_INCLUDE_DIRS})
+# to generate trt and test image dir
+add_executable(yolov5 ${PROJECT_SOURCE_DIR}/yolov5.cpp)
+target_link_libraries(yolov5 nvinfer cudart myplugins ${OpenCV_LIBS})
+add_definitions(-O3 -pthread)
 
-add_executable(yolop ${PROJECT_SOURCE_DIR}/main.cpp)
-target_link_libraries(yolop nvinfer)
-target_link_libraries(yolop ${ZED_LIBS})
-target_link_libraries(yolop cudart)
-target_link_libraries(yolop myplugins)
-target_link_libraries(yolop ${OpenCV_LIBS})
 
-add_definitions(-O3 -pthread)
+# to test with zed camera
+if(ZED_FOUND)
+    set(ZED_LIBS ${ZED_LIBRARIES} ${CUDA_CUDA_LIBRARY} ${CUDA_CUDART_LIBRARY})
 
+    add_executable(yolop ${PROJECT_SOURCE_DIR}/main.cpp)
+    target_link_libraries(yolop nvinfer)
+    target_link_libraries(yolop ${ZED_LIBS})
+    target_link_libraries(yolop cudart)
+    target_link_libraries(yolop myplugins)
+    target_link_libraries(yolop ${OpenCV_LIBS})
+endif(ZED_FOUND)

toolkits/deploy/README.md

@@ -0,0 +1,63 @@
+YoloP TensorRT Usage(简要说明)
+=====
+
+
+## 1. 准备构建环境 | Prepare building environments
+
+Make sure you have install `c++`(support c++11)、 `cmake`、`opencv`(4.x)、`cuda`(10.x)、`nvinfer`(7.x).
+
+Now, zedcam is not necessary!
+
+And we can also use opencv that built without cuda.
+
+## 2. 编译 | build
+
+Go to `YOLOP/toolkits/deploy`.
+
+```
+mkdir build
+cd build
+
+cmake ..
+make
+```
+
+Now you can get `yolov5` and `libmyplugins.so`.
+If you have Zed installed, you can get `yolop`.
+
+
+## 3. 生成和测试 trt | Generate and test trt
+
+Go to build dir (`YOLOP/toolkits/deploy/build`).
+
+### 3.1 gen wts
+```
+python3 ../gen_wts.py
+```
+
+### 3.2 gen trt
+```
+./yolov5 -s yolop.wts  yolop.trt s
+```
+
+### 3.3 test trt
+```
+mkdir ../results
+./yolov5 -d yolop.trt  ../../../inference/images/
+```
+
+It will output like as follow if successful! (`Jetson Xavier NX - Jetpack 4.4`)
+```
+1601ms
+26ms
+26ms
+26ms
+26ms
+28ms
+```
+
+![](build/results/_3c0e7240-96e390d2.jpg)
+
+
+
+

toolkits/deploy/gen_wts.py

@@ -11,7 +11,7 @@
 device = torch.device('cpu')
 # Load model
 model = get_net(cfg)
-checkpoint = torch.load('weights/End-to-end.pth', map_location=device)
+checkpoint = torch.load(BASE_DIR + '/weights/End-to-end.pth', map_location=device)
 model.load_state_dict(checkpoint['state_dict'])
 # load to FP32
 model.float()