ソースを参照

add trafficlight_config.yaml

fujiankuan 1 年間 前
コミット
43f1134187

+ 65 - 23
src/detection/trafficlight_detection/main.cpp

@@ -55,19 +55,21 @@ struct Bbox {
 };
 
 
-const char* onnx_path = "./model/yolov5s_640_old+new.onnx"; //onnx path
-const std::string engine_path = "./model/yolov5s_640_old+new.engine"; // engine paths
+string config_file = "./yaml/trafficlight_config.yaml";
+string onnx_path = "./model/yolov5s_640_old+new.onnx"; //onnx path
+string engine_path = "./model/yolov5s_640_old+new.engine"; // engine paths
 
 bool imshow_flag = true;
+bool ivlog_flag = false;
 
 float conf_thr = 0.5;
 float nms_thr = 0.4;
 
-const bool calibrationstate = false;
-const std::string calibration_yamlpath = "./yaml/camera_middle_640_360.yaml";
+bool calibrationstate = false;
+string calibration_yamlpath = "./yaml/camera_middle_640_360.yaml";
 
 
-const bool cropstate = true;
+bool cropstate = true;
 cv::Range crop_height = cv::Range(40,680);
 cv::Range crop_width = cv::Range(320,960);
 
@@ -77,7 +79,7 @@ bool test_video = false;
 string video_path = "/home/nvidia/code/modularization/src/detection/trafficlight_detection/data/camera_test3.mp4";
 
 //是否开启跟踪
-const bool trackstate = true;
+bool trackstate = true;
 
 void * gpbrain;  //获取红绿灯开启状态
 string brainname = "brainstate";
@@ -145,7 +147,7 @@ ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, IBuilder
     INetworkDefinition* network = builder->createNetworkV2(1U); //此处重点1U为OU就有问题
 
     IParser* parser = createParser(*network, gLogger);
-    parser->parseFromFile(onnx_path, static_cast<int32_t>(ILogger::Severity::kWARNING));
+    parser->parseFromFile(onnx_path.c_str(), static_cast<int32_t>(ILogger::Severity::kWARNING));
     //解析有错误将返回
     for (int32_t i = 0; i < parser->getNbErrors(); ++i) { std::cout << parser->getError(i)->desc() << std::endl; }
     std::cout << "successfully parse the onnx model" << std::endl;
@@ -562,15 +564,16 @@ void shareLightMsg(vector<Detection> results)
         light->set_index(i+1);
 
         int light_type = 255;
-        if(results.at(i).class_id == 0)
+        if (results.at(i).class_id == 0)  //红色
             light_type = 2;
-        else if(results.at(i).class_id == 1)
+        else if(results.at(i).class_id == 1)  //绿色
             light_type = 1;
-        else if(results.at(i).class_id == 2)
+        else if(results.at(i).class_id == 2)  //黄色
             light_type = 3;
         light->set_type(light_type);
-        //light->set_type(results.at(i).class_id+1);
-        cout<<"#######################################: "<<light_type<<endl;
+        cout<<"light color: "<<light_type<<endl;
+        if(ivlog_flag)
+            givlog->verbose("light color: %d",light_type);
     }
 
     int size = light_array.ByteSize();
@@ -716,16 +719,42 @@ int main(int argc, char** argv)
     givlog = new iv::Ivlog("traffic_light_detection");
     gfault->SetFaultState(0,0,"yolov5 initialize.");
 
-    if(argc==3)
-    {
-        test_video = (strcmp(argv[1], "true") == 0)?true:false;
-        video_path = argv[2];
-    }
-    if(argc==2)
+//    if(argc==3)
+//    {
+//        test_video = (strcmp(argv[1], "true") == 0)?true:false;
+//        video_path = argv[2];
+//    }
+//    if(argc==2)
+//    {
+//        test_video = (strcmp(argv[1], "true") == 0)?true:false;
+//    }
+
+
+    cv::FileStorage config(config_file, cv::FileStorage::READ);
+    bool config_isOpened = config.isOpened();
+    //const char* onnx_path_;
+    if(config_isOpened)
     {
-        test_video = (strcmp(argv[1], "true") == 0)?true:false;
+        onnx_path = string(config["model"]["onnx_path"]);
+        engine_path = string(config["model"]["engine_path"]);
+        imshow_flag = (string(config["imshow_flag"]) == "true");
+        ivlog_flag = (string(config["ivlog_flag"]) == "true");
+        test_video = (string(config["test_video"]["open"]) == "true");
+        video_path = string(config["test_video"]["video_path"]);
+        lightstart = (string(config["lightstart"]) == "true");
+        trackstate = (string(config["trackstate"]) == "true");
+        conf_thr = float(config["conf_thr"]);
+        nms_thr = float(config["nms_thr"]);
+        calibrationstate = (string(config["camera"]["calibrationstate"]) == "true");
+        calibration_yamlpath = string(config["camera"]["calibration_yamlpath"]);
+        cropstate = (string(config["camera"]["cropstate"]) == "true");
+        crop_height = cv::Range(int(config["camera"]["crop_height"]["start"]),
+                                int(config["camera"]["crop_height"]["end"]));
+
+        crop_width = cv::Range(int(config["camera"]["crop_width"]["start"]),
+                               int(config["camera"]["crop_width"]["end"]));
     }
-
+    config.release();
     if(test_video)
         std::thread * readthread = new std::thread(ReadFunc,1);
     else
@@ -767,7 +796,7 @@ int main(int argc, char** argv)
                 {
                     double waittotal = (double)cv::getTickCount() - waittime;
                     double totaltime = waittotal/cv::getTickFrequency();
-//                    if(totaltime>120.0)
+//                    if(totaltime>10.0)
 //                    {
 //                        cout<<"Cant't get frame and quit"<<endl;
 //                        lightstart = false;
@@ -776,6 +805,9 @@ int main(int argc, char** argv)
 //                        break;
 //                    }
                     cout<<"Wait for frame "<<totaltime<<"s"<<endl;
+                    if(ivlog_flag)
+                        givlog->verbose("Wait for frame %f s",totaltime);
+                    std::this_thread::sleep_for(std::chrono::milliseconds(10));
                     continue;
                 }
                 auto start = std::chrono::system_clock::now();  //时间函数
@@ -797,6 +829,8 @@ int main(int argc, char** argv)
                     vector<od::bbox_t> outs_90;
                     for (int i = 0; i < results.size(); i++)
                     {
+                        //-------------判断红绿灯是否为横向,width=(x1-x2),height=(y1-y2)-----------
+
                         bbox_t_90.x = results.at(i).bbox[0];
                         bbox_t_90.y = results.at(i).bbox[1];
                         bbox_t_90.w = results.at(i).bbox[2];
@@ -851,11 +885,13 @@ int main(int argc, char** argv)
                 auto end = std::chrono::system_clock::now();  //时间函数
                 std::cout <<"total time traffic: "<< std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
 
-                //显示检测检测结果
+                //显示检测结果,检测结果没有返回到原始图像大小,所以需要resize到网络输入尺寸
                 cv::resize(frame, frame, cv::Size(INPUT_W, INPUT_H), 0, 0, cv::INTER_LINEAR);
                 res_img=draw_rect(frame,results_final);
 
                 std::cout<<"light number: "<<results_final.size()<<std::endl;
+                if(ivlog_flag)
+                    givlog->verbose("light number: %d",results_final.size());
 
                 results.clear();
                 results_track.clear();
@@ -889,7 +925,13 @@ int main(int argc, char** argv)
                 }
             }
         }
-        //else cout<<"Wait for lightstart==true"<<endl;
+        else
+        {
+            cout<<"Wait for lightstart==true"<<endl;
+            if(ivlog_flag)
+                givlog->verbose("Wait for lightstart==true");
+            std::this_thread::sleep_for(std::chrono::milliseconds(10));
+        }
     }
 
     // Destroy the engine

+ 27 - 0
src/detection/trafficlight_detection/yaml/trafficlight_config.yaml

@@ -0,0 +1,27 @@
+%YAML:1.0
+---
+imshow_flag: true
+ivlog_flag: false
+test_video:
+  open: false
+  video_path: /home/nvidia/code/modularization/src/detection/trafficlight_detection/data/camera_test2.mp4
+lightstart: true 
+trackstate: true
+conf_thr: 0.5
+nms_thr : 0.4
+model:
+  onnx_path: ./model/yolov5s_640_old+new.onnx
+  engine_path: ./model/yolov5s_640_old+new.engine
+camera:
+  calibrationstate : false
+  calibration_yamlpath : ./yaml/camera_middle_640_360.yaml
+  cropstate: true
+  crop_height: 
+    start: 40
+    end: 680
+  crop_width:
+    start: 320
+    end: 960
+
+
+