1

1

Author: LittleSheepy

tensorrtx_my/yolov5/CMakeLists.txt

@@ -53,15 +53,15 @@ ELSEIF(CMAKE_SYSTEM_NAME MATCHES "Windows")
 	set(CUDA_NVCC_PLAGS ${CUDA_NVCC_PLAGS};-std=c++11; -g; -G;-gencode; arch=compute_75;code=sm_75)
 	find_package(CUDA REQUIRED)
 	include_directories(${CUDA_INCLUDE_DIRS})
-	enable_language(CUDA)  # add this line, then no need to setup cuda path in vs
+	#enable_language(CUDA)  # add this line, then no need to setup cuda path in vs
 
 	# tensorrt
-	set(TRT_DIR "C://tools//TensorRT-7.0.0.11//")  #3
+	set(TRT_DIR "D:/SDK/TensorRT-8.2.0.6//")  #3
 	include_directories(${TRT_DIR}\\include)
 	link_directories(${TRT_DIR}\\lib)
 
 	# setup opencv
-	set(OpenCV_DIR "C:/tools/opencv/build/")  #2
+	set(OpenCV_DIR "D:/sheepy/tools/opencv455/build/")  #2
     find_package(OpenCV REQUIRED)
 	include_directories(${OpenCV_INCLUDE_DIRS})
 	# Threads

tensorrtx_my/yolov5_7.0/CMakeLists.txt

@@ -54,11 +54,11 @@ ELSEIF(CMAKE_SYSTEM_NAME MATCHES "Windows")
 	set(CUDA_NVCC_PLAGS ${CUDA_NVCC_PLAGS};-std=c++11; -g; -G;-gencode; arch=compute_75;code=sm_75)
 	find_package(CUDA REQUIRED)
 	include_directories(${CUDA_INCLUDE_DIRS})
-	link_directories("C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v11.3/lib/x64/")
+	link_directories("C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v11.1/lib/x64/")
 	enable_language(CUDA)  # add this line, then no need to setup cuda path in vs
 
 	# tensorrt
-	set(TRT_DIR "D://tools//TensorRT-8.0.1.6-cu11.3//")  #3
+	set(TRT_DIR "D:/tools/TensorRT-8.2.0.6//")  #3
 	include_directories(${TRT_DIR}\\include)
 	link_directories(${TRT_DIR}\\lib)
 
@@ -104,7 +104,7 @@ ELSEIF(CMAKE_SYSTEM_NAME MATCHES "Windows")
 ENDIF()
 
 #生成目标文件
-add_executable(${PROJECT_NAME} ${PROJECT_SOURCE_DIR}/yolov5_det.cpp ${SRC_FILES} ${SRC_CU_FILES} ${INC_FILES})
+add_executable(${PROJECT_NAME} ${PROJECT_SOURCE_DIR}/yolov5_seg.cpp ${SRC_FILES} ${SRC_CU_FILES} ${INC_FILES})
 
 #target_link_libraries(${PROJECT_NAME} ${LIB_NAME})
 target_link_libraries(${PROJECT_NAME} nvinfer cudart nvinfer_plugin ${CUDA_LIBRARIES})

tensorrtx_my/yolov5_7.0/include/config.h

@@ -23,8 +23,8 @@ constexpr static int kClsNumClass = 1000;
 constexpr static int kBatchSize = 1;
 
 // Yolo's input width and height must by divisible by 32
-constexpr static int kInputH = 320;
-constexpr static int kInputW = 320;
+constexpr static int kInputH = 1024;
+constexpr static int kInputW = 1024;
 
 // Classfication model's input shape
 constexpr static int kClsInputH = 224;
@@ -46,7 +46,7 @@ constexpr static float kIgnoreThresh = 0.1f;
 
 // NMS overlapping thresh and final detection confidence thresh
 const static float kNmsThresh = 0.45f;
-const static float kConfThresh = 0.5f;
+const static float kConfThresh = 0.6f;
 
 const static int kGpuId = 0;
 

tensorrtx_my/yolov5_7.0/seg_best.wts

@@ -24,6 +24,15 @@ cv::Rect get_rect(cv::Mat& img, float bbox[4]) {
     t = t / r_h;
     b = b / r_h;
   }
+  if (l < 0) {
+      l = 0;
+  }
+  if (b >= 2048) {
+      b = 2048;
+  }
+  if (r >= 2448) {
+      r = 2448;
+  }
   return cv::Rect(round(l), round(t), round(r - l), round(b - t));
 }
 
@@ -93,6 +102,9 @@ void draw_bbox(std::vector<cv::Mat>& img_batch, std::vector<std::vector<Detectio
 
 static cv::Rect get_downscale_rect(float bbox[4], float scale) {
   float left = bbox[0] - bbox[2] / 2;
+  if (left < 0) {
+      left = 0;
+  }
   float top = bbox[1] - bbox[3] / 2;
   float right = bbox[0] + bbox[2] / 2;
   float bottom = bbox[1] + bbox[3] / 2;
@@ -110,11 +122,17 @@ std::vector<cv::Mat> process_mask(const float* proto, int proto_size, std::vecto
     auto r = get_downscale_rect(dets[i].bbox, 4);
     for (int x = r.x; x < r.x + r.width; x++) {
       for (int y = r.y; y < r.y + r.height; y++) {
+          if ((x<=0)||(y<=0)) {
+              continue;
+          }
         float e = 0.0f;
         for (int j = 0; j < 32; j++) {
           e += dets[i].mask[j] * proto[j * proto_size / 32 + y * mask_mat.cols + x];
         }
         e = 1.0f / (1.0f + expf(-e));
+        if ((x>= kInputW / 4)||(y>= kInputH / 4)) {
+            continue;
+        }
         mask_mat.at<float>(y, x) = e;
       }
     }
@@ -158,6 +176,9 @@ void draw_mask_bbox(cv::Mat& img, std::vector<Detection>& dets, std::vector<cv::
     cv::Rect r = get_rect(img, dets[i].bbox);
     for (int x = r.x; x < r.x + r.width; x++) {
       for (int y = r.y; y < r.y + r.height; y++) {
+          if ((x <= 0) || (y <= 0)) {
+              continue;
+          }
         float val = img_mask.at<float>(y, x);
         if (val <= 0.5) continue;
         img.at<cv::Vec3b>(y, x)[0] = img.at<cv::Vec3b>(y, x)[0] / 2 + bgr[0] / 2;

tensorrtx_my/yolov5_7.0/src/postprocess.cpp

@@ -194,11 +194,12 @@ int main(int argc, char** argv) {
 
   // batch predict
   for (size_t i = 0; i < file_names.size(); i += kBatchSize) {
+    std::cout << "i:" << i << std::endl;
     // Get a batch of images
     std::vector<cv::Mat> img_batch;
     std::vector<std::string> img_name_batch;
     for (size_t j = i; j < i + kBatchSize && j < file_names.size(); j++) {
-      cv::Mat img = cv::imread(img_dir + "/" + file_names[j]);
+        cv::Mat img = cv::imread(img_dir + "/" + file_names[j]);
       img_batch.push_back(img);
       img_name_batch.push_back(file_names[j]);
     }
@@ -220,10 +221,15 @@ int main(int argc, char** argv) {
     for (size_t b = 0; b < img_name_batch.size(); b++) {
       auto& res = res_batch[b];
       cv::Mat img = img_batch[b];
+      cv::Mat img_org;
+      img_batch[b].copyTo(img_org);
 
       auto masks = process_mask(&cpu_output_buffer2[b * kOutputSize2], kOutputSize2, res);
       draw_mask_bbox(img, res, masks, labels_map);
-      cv::imwrite("_" + img_name_batch[b], img);
+      if (res.size() > 0) {
+          cv::imwrite("result/" + img_name_batch[b], img);
+          cv::imwrite("NG/" + img_name_batch[b], img_org);
+      }
     }
   }