refactor yolov5 dir

README.md

@@ -15,7 +15,7 @@ The basic workflow of TensorRTx is:
 
 ## News
 
-- `18 Dec 2022`. [YOLOv5](./yolov5) upgrade to support v7.0, including instance segmention.
+- `18 Dec 2022`. [YOLOv5](./yolov5) upgrade to support v7.0, including instance segmentation.
 - `12 Dec 2022`. [East-Face](https://github.com/East-Face): [UNet](./unet) upgrade to support v3.0 of [Pytorch-UNet](https://github.com/milesial/Pytorch-UNet).
 - `26 Oct 2022`. [ausk](https://github.com/ausk): YoloP(You Only Look Once for Panopitic Driving Perception).
 - `19 Sep 2022`. [QIANXUNZDL123](https://github.com/QIANXUNZDL123) and [lindsayshuo](https://github.com/lindsayshuo): YOLOv7.

yolov5/CMakeLists.txt

@@ -1,4 +1,4 @@
-cmake_minimum_required(VERSION 2.6)
+cmake_minimum_required(VERSION 3.10)
 
 project(yolov5)
 
@@ -8,51 +8,45 @@ option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
 set(CMAKE_CXX_STANDARD 11)
 set(CMAKE_BUILD_TYPE Debug)
 
-find_package(CUDA REQUIRED)
-
-if(WIN32)
+# TODO(Call for PR): make cmake compatible with Windows
+set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc)
 enable_language(CUDA)
-endif(WIN32)
 
-include_directories(${PROJECT_SOURCE_DIR}/include)
 # include and link dirs of cuda and tensorrt, you need adapt them if yours are different
 # cuda
 include_directories(/usr/local/cuda/include)
 link_directories(/usr/local/cuda/lib64)
 # tensorrt
 # TODO(Call for PR): make TRT path configurable from command line
-include_directories(/usr/include/x86_64-linux-gnu/)
-link_directories(/usr/lib/x86_64-linux-gnu/)
+include_directories(/home/nvidia/TensorRT-8.2.5.1/include/)
+link_directories(/home/nvidia/TensorRT-8.2.5.1/lib/)
 
-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -Ofast -g -Wfatal-errors -D_MWAITXINTRIN_H_INCLUDED")
-cuda_add_library(myplugins SHARED yololayer.cu)
+include_directories(${PROJECT_SOURCE_DIR}/src/)
+include_directories(${PROJECT_SOURCE_DIR}/plugin/)
+file(GLOB_RECURSE SRCS ${PROJECT_SOURCE_DIR}/src/*.cpp ${PROJECT_SOURCE_DIR}/src/*.cu)
+file(GLOB_RECURSE PLUGIN_SRCS ${PROJECT_SOURCE_DIR}/plugin/*.cu)
+
+add_library(myplugins SHARED ${PLUGIN_SRCS})
 target_link_libraries(myplugins nvinfer cudart)
 
 find_package(OpenCV)
 include_directories(${OpenCV_INCLUDE_DIRS})
 
-cuda_add_executable(yolov5 calibrator.cpp yolov5.cpp preprocess.cu)
+add_executable(yolov5_det yolov5_det.cpp ${SRCS})
+target_link_libraries(yolov5_det nvinfer)
+target_link_libraries(yolov5_det cudart)
+target_link_libraries(yolov5_det myplugins)
+target_link_libraries(yolov5_det ${OpenCV_LIBS})
 
-target_link_libraries(yolov5 nvinfer)
-target_link_libraries(yolov5 cudart)
-target_link_libraries(yolov5 myplugins)
-target_link_libraries(yolov5 ${OpenCV_LIBS})
+add_executable(yolov5_cls yolov5_cls.cpp ${SRCS})
+target_link_libraries(yolov5_cls nvinfer)
+target_link_libraries(yolov5_cls cudart)
+target_link_libraries(yolov5_cls myplugins)
+target_link_libraries(yolov5_cls ${OpenCV_LIBS})
 
-add_executable(yolov5-cls calibrator.cpp yolov5_cls.cpp)
-
-target_link_libraries(yolov5-cls nvinfer)
-target_link_libraries(yolov5-cls cudart)
-target_link_libraries(yolov5-cls myplugins)
-target_link_libraries(yolov5-cls ${OpenCV_LIBS})
-
-cuda_add_executable(yolov5-seg calibrator.cpp yolov5_seg.cpp preprocess.cu)
-
-target_link_libraries(yolov5-seg nvinfer)
-target_link_libraries(yolov5-seg cudart)
-target_link_libraries(yolov5-seg myplugins)
-target_link_libraries(yolov5-seg ${OpenCV_LIBS})
-
-if(UNIX)
-add_definitions(-O2 -pthread)
-endif(UNIX)
+add_executable(yolov5_seg yolov5_seg.cpp ${SRCS})
+target_link_libraries(yolov5_seg nvinfer)
+target_link_libraries(yolov5_seg cudart)
+target_link_libraries(yolov5_seg myplugins)
+target_link_libraries(yolov5_seg ${OpenCV_LIBS})
 

yolov5/README.md

@@ -81,14 +81,14 @@ cd build
 cp {ultralytics}/yolov5/yolov5s.wts {tensorrtx}/yolov5/build
 cmake ..
 make
-sudo ./yolov5 -s [.wts] [.engine] [n/s/m/l/x/n6/s6/m6/l6/x6 or c/c6 gd gw]  // serialize model to plan file
-sudo ./yolov5 -d [.engine] [image folder]  // deserialize and run inference, the images in [image folder] will be processed.
+sudo ./yolov5_det -s [.wts] [.engine] [n/s/m/l/x/n6/s6/m6/l6/x6 or c/c6 gd gw]  // serialize model to plan file
+sudo ./yolov5_det -d [.engine] [image folder]  // deserialize and run inference, the images in [image folder] will be processed.
 // For example yolov5s
-sudo ./yolov5 -s yolov5s.wts yolov5s.engine s
-sudo ./yolov5 -d yolov5s.engine ../samples
+sudo ./yolov5_det -s yolov5s.wts yolov5s.engine s
+sudo ./yolov5_det -d yolov5s.engine ../samples
 // For example Custom model with depth_multiple=0.17, width_multiple=0.25 in yolov5.yaml
-sudo ./yolov5 -s yolov5_custom.wts yolov5.engine c 0.17 0.25
-sudo ./yolov5 -d yolov5.engine ../samples
+sudo ./yolov5_det -s yolov5_custom.wts yolov5.engine c 0.17 0.25
+sudo ./yolov5_det -d yolov5.engine ../samples
 ```
 
 3. check the images generated, as follows. _zidane.jpg and _bus.jpg
@@ -98,10 +98,10 @@ sudo ./yolov5 -d yolov5.engine ../samples
 ```
 // install python-tensorrt, pycuda, etc.
 // ensure the yolov5s.engine and libmyplugins.so have been built
-python yolov5_trt.py
+python yolov5_det_trt.py
 
 // Another version of python script, which is using CUDA Python instead of pycuda.
-python yolov5_trt_cuda_python.py
+python yolov5_det_trt_cuda_python.py
 ```
 
 <p align="center">
@@ -115,20 +115,20 @@ python yolov5_trt_cuda_python.py
 wget https://github.com/joannzhang00/ImageNet-dataset-classes-labels/blob/main/imagenet_classes.txt
 
 # Build and serialize TensorRT engine
-./yolov5-cls -s yolov5s-cls.wts yolov5s-cls.engine s
+./yolov5_cls -s yolov5s-cls.wts yolov5s-cls.engine s
 
 # Run inference
-./yolov5-cls -d yolov5s-cls.engine ../samples
+./yolov5_cls -d yolov5s-cls.engine ../samples
 ```
 
 ### Instance Segmentation
 
 ```
 # Build and serialize TensorRT engine
-./yolov5-seg -s yolov5s-seg.wts yolov5s-seg.engine s
+./yolov5_seg -s yolov5s-seg.wts yolov5s-seg.engine s
 
 # Run inference
-./yolov5-seg -d yolov5s-seg.engine ../samples
+./yolov5_seg -d yolov5s-seg.engine ../samples
 ```
 
 <p align="center">

yolov5/gen_wts.py

@@ -30,11 +30,10 @@ def parse_args():
 pt_file, wts_file, m_type = parse_args()
 print(f'Generating .wts for {m_type} model')
 
-# Initialize
-device = select_device('cpu')
 # Load model
 print(f'Loading {pt_file}')
-model = torch.load(pt_file, map_location=device)  # load to FP32
+device = select_device('cpu')
+model = torch.load(pt_file, map_location=device)  # Load FP32 weights
 model = model['ema' if model.get('ema') else 'model'].float()
 
 if m_type in ['detect', 'seg']:

yolov5/yolov5_cls.cpp

@@ -212,8 +212,8 @@ int main(int argc, char** argv) {
     std::string img_dir;
     if (!parse_args(argc, argv, wts_name, engine_name, gd, gw, img_dir)) {
         std::cerr << "arguments not right!" << std::endl;
-        std::cerr << "./yolov5-cls -s [.wts] [.engine] [n/s/m/l/x or c gd gw]  // serialize model to plan file" << std::endl;
-        std::cerr << "./yolov5-cls -d [.engine] ../samples  // deserialize plan file and run inference" << std::endl;
+        std::cerr << "./yolov5_cls -s [.wts] [.engine] [n/s/m/l/x or c gd gw]  // serialize model to plan file" << std::endl;
+        std::cerr << "./yolov5_cls -d [.engine] ../samples  // deserialize plan file and run inference" << std::endl;
         return -1;
     }
 

yolov5/yolov5_cls_trt.py

@@ -216,7 +216,7 @@ class warmUpThread(threading.Thread):
 
 if __name__ == "__main__":
     # load custom plugin and engine
-    engine_file_path = "build/yolov5s_cls.engine"
+    engine_file_path = "build/yolov5s-cls.engine"
 
     if len(sys.argv) > 1:
         engine_file_path = sys.argv[1]

yolov5/yolov5_det.cpp

@@ -305,8 +305,8 @@ int main(int argc, char** argv) {
     std::string img_dir;
     if (!parse_args(argc, argv, wts_name, engine_name, is_p6, gd, gw, img_dir)) {
         std::cerr << "arguments not right!" << std::endl;
-        std::cerr << "./yolov5 -s [.wts] [.engine] [n/s/m/l/x/n6/s6/m6/l6/x6 or c/c6 gd gw]  // serialize model to plan file" << std::endl;
-        std::cerr << "./yolov5 -d [.engine] ../samples  // deserialize plan file and run inference" << std::endl;
+        std::cerr << "./yolov5_det -s [.wts] [.engine] [n/s/m/l/x/n6/s6/m6/l6/x6 or c/c6 gd gw]  // serialize model to plan file" << std::endl;
+        std::cerr << "./yolov5_det -d [.engine] ../samples  // deserialize plan file and run inference" << std::endl;
         return -1;
     }
 

yolov5/yolov5_seg.cpp

@@ -287,8 +287,8 @@ int main(int argc, char** argv) {
     std::string img_dir;
     if (!parse_args(argc, argv, wts_name, engine_name, gd, gw, img_dir)) {
         std::cerr << "arguments not right!" << std::endl;
-        std::cerr << "./yolov5-seg -s [.wts] [.engine] [n/s/m/l/x or c gd gw]  // serialize model to plan file" << std::endl;
-        std::cerr << "./yolov5-seg -d [.engine] ../samples  // deserialize plan file and run inference" << std::endl;
+        std::cerr << "./yolov5_seg -s [.wts] [.engine] [n/s/m/l/x or c gd gw]  // serialize model to plan file" << std::endl;
+        std::cerr << "./yolov5_seg -d [.engine] ../samples  // deserialize plan file and run inference" << std::endl;
         return -1;
     }