fix problem about tensorrt

.gitignore

@@ -63,5 +63,5 @@
 *.dSYM/
 
 
-*.pb.cc
-*.pb.h
+*.pb.*
+*.pb.*

CMakeLists.txt

@@ -3,7 +3,7 @@ set(SON_PROJECT_NAME tof3d)
 message("========== Load son project ${SON_PROJECT_NAME} ==========" )
 
 unset(OPTION_ENABLE_TENSORRT_DETECT_CODE CACHE)
-option(OPTION_ENABLE_TENSORRT_DETECT_CODE "" OFF)
+option(OPTION_ENABLE_TENSORRT_DETECT_CODE "" ON)
 message("<=${SON_PROJECT_NAME}=> OPTION_ENABLE_TENSORRT_DETECT_CODE: " ${OPTION_ENABLE_TENSORRT_DETECT_CODE})
 
 if (OPTION_ENABLE_TENSORRT_DETECT_CODE)

detect/onnx/wheel-detector.h

@@ -2,15 +2,15 @@
 
 #include "inference.h"
 
+struct Object {
+    cv::Rect_<float> rect;
+    int              label = 0;
+    float            prob  = 0.0;
+    cv::Mat          boxMask;
+};
+    
 class TensorrtWheelDetector{
 public:
-    struct Object {
-        cv::Rect_<float> rect;
-        int              label = 0;
-        float            prob  = 0.0;
-        cv::Mat          boxMask;
-    };
-
     TensorrtWheelDetector(const std::string &model_file, const std::string &class_file);
     ~TensorrtWheelDetector();
 

detect/tensorrt/wheel-detector.cpp

@@ -1,7 +1,9 @@
 #include "wheel-detector.h"
 
-TensorrtWheelDetector::TensorrtWheelDetector(const std::string &enginfile){
-    yolov8_=new YOLOv8_seg(enginfile);
+TensorrtWheelDetector::TensorrtWheelDetector(const std::string &model_file, const std::string &class_file){
+    cudaSetDevice(0);
+    
+    yolov8_ = new YOLOv8_seg(model_file);
     yolov8_->make_pipe(false);
     imgsz_=cv::Size{640, 480};
     seg_h_        = 120;

detect/tensorrt/wheel-detector.h

@@ -4,7 +4,7 @@
 
 class TensorrtWheelDetector{
 public:
-    TensorrtWheelDetector(const std::string &enginfile);
+    TensorrtWheelDetector(const std::string &model_file, const std::string &class_file);
     ~TensorrtWheelDetector();
 
 
@@ -22,4 +22,4 @@ private:
 
 };
 
-#endif
+#endif

vzense/device_tof3d.cpp

@@ -7,8 +7,11 @@ Error_manager DeviceTof3D::Init(const VzEtcMap &tp_tof3d, bool search) {
         return {TOF3D_VZENSE_DEVICE_INIT_FAILED, MAJOR_ERROR, "VzInitialize failed status: %d", m_vz_status};
     }
 
+#ifdef ENABLE_TENSORRT_DETECT
+    detector = new TensorrtWheelDetector(ETC_PATH PROJECT_NAME "/model.engine", ETC_PATH PROJECT_NAME "/class.txt");
+#else
     detector = new TensorrtWheelDetector(ETC_PATH PROJECT_NAME "/model.onnx", ETC_PATH PROJECT_NAME "/class.txt");
-
+#endif
     loadEtc(tp_tof3d);
 
     if (search) {
@@ -449,7 +452,7 @@ void DeviceTof3D::receiveFrameThread(const std::string &ip) {
 
     pcl::PointCloud<pcl::PointXYZ>::Ptr t_car_cloud(new pcl::PointCloud<pcl::PointXYZ>);
     pcl::PointCloud<pcl::PointXYZ>::Ptr t_wheel_cloud(new pcl::PointCloud<pcl::PointXYZ>);
-    std::vector<TensorrtWheelDetector::Object> segment_results;
+    std::vector<Object> segment_results;
     while (t_iter->second.condit->is_alive()) {
         t_iter->second.condit->wait();
         if (t_iter->second.condit->is_alive()) {
@@ -463,8 +466,8 @@ void DeviceTof3D::receiveFrameThread(const std::string &ip) {
                     segment_results.clear();
                     if (Frame2Mat(depthFrame, depthMat) == SUCCESS) {
                         detector->detect(depthMat, segment_results);
-                        cv::imshow("ret", depthMat);
-                        cv::waitKey(100);
+                        // cv::imshow("ret", depthMat);
+                        // cv::waitKey(100);
                     }
 
                     std::vector<cv::Point> mat_seg_ret;