main.cpp 32 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948
  1. #include "NvInfer.h"
  2. #include "cuda_runtime_api.h"
  3. #include <fstream>
  4. #include <iostream>
  5. #include <map>
  6. #include <sstream>
  7. #include <vector>
  8. #include <chrono>
  9. #include <cmath>
  10. #include <cassert>
  11. #include <algorithm>
  12. #include<opencv2/core/core.hpp>
  13. #include<opencv2/highgui/highgui.hpp>
  14. #include <opencv2/opencv.hpp>
  15. #include "modulecomm.h"
  16. #include "xmlparam.h"
  17. #include "ivfault.h"
  18. #include "ivlog.h"
  19. #include "ivexit.h"
  20. #include "ivversion.h"
  21. #include "rawpic.pb.h"
  22. #include "lightarray.pb.h"
  23. #include "brainstate.pb.h"
  24. #include <QCoreApplication>
  25. #include <thread>
  26. #include "opencv2/imgcodecs/legacy/constants_c.h"
  27. #include "qmutex.h"
  28. #include "condition_variable"
  29. #include "imageBuffer.h"
  30. #include "detect_obstacle.h"
  31. // onnx转换头文件
  32. #include "NvOnnxParser.h"
  33. using namespace nvonnxparser;
  34. using namespace nvinfer1;
  35. using namespace std;
  36. //全局变量
  37. struct Detection {
  38. //center_x center_y w h
  39. float bbox[4];
  40. float conf; // bbox_conf * cls_conf
  41. int class_id;
  42. int index;
  43. };
  44. struct Bbox {
  45. int x;
  46. int y;
  47. int w;
  48. int h;
  49. };
  50. string config_file = "./yaml/trafficlight_config.yaml";
  51. string onnx_path = "./model/yolov5s_640_old+new.onnx"; //onnx path
  52. string engine_path = "./model/yolov5s_640_old+new.engine"; // engine paths
  53. bool imshow_flag = true;
  54. bool ivlog_flag = false;
  55. float conf_thr = 0.5;
  56. float nms_thr = 0.4;
  57. bool calibrationstate = false;
  58. string calibration_yamlpath = "./yaml/camera_middle_640_360.yaml";
  59. bool cropstate = true;
  60. cv::Range crop_height = cv::Range(40,680);
  61. cv::Range crop_width = cv::Range(320,960);
  62. bool test_video = false;
  63. //string video_path = "v4l2src device=/dev/video0 ! video/x-raw, width=(int)1280, height=(int)720 ! videoconvert ! appsink";
  64. string video_path = "/home/nvidia/code/modularization/src/detection/trafficlight_detection/data/camera_test3.mp4";
  65. //是否开启跟踪
  66. bool trackstate = true;
  67. void * gpbrain; //获取红绿灯开启状态
  68. string brainname = "brainstate";
  69. bool lightstart = false;
  70. void * gpcamera;
  71. string cameraname="image00";
  72. ConsumerProducerQueue<cv::Mat> * imageBuffer = new ConsumerProducerQueue<cv::Mat>(3,true);
  73. void * gpdetect;
  74. string detectname = "lightarray"; //检测结果
  75. // stuff we know about the network and the input/output blobs
  76. static const int INPUT_H = 640;
  77. static const int INPUT_W = 640;
  78. static const int cls_num = 3;
  79. //不同输入尺寸anchor:960-->56700 | 640-->25200 | 416-->10647 | 320-->6300
  80. static const int anchor_output_num = 25200;
  81. static const int OUTPUT_SIZE = 1* anchor_output_num *(cls_num+5); //1000 * sizeof(Detection) / sizeof(float) + 1;
  82. const char* INPUT_BLOB_NAME = "images";
  83. const char* OUTPUT_BLOB_NAME = "output";
  84. IExecutionContext* context;
  85. ICudaEngine* engine;
  86. IRuntime* runtime;
  87. float* prob;
  88. iv::Ivfault *gfault = nullptr;
  89. iv::Ivlog *givlog = nullptr;
  90. cv::Mat camera_matrix,dist_coe,map1,map2; //标定参数
  91. float time_read_img = 0.0;
  92. float time_infer = 0.0;
  93. int time_num = 0;
  94. #define CHECK(status) \
  95. do\
  96. {\
  97. auto ret = (status);\
  98. if (ret != 0)\
  99. {\
  100. std::cerr << "Cuda failure: " << ret << std::endl;\
  101. abort();\
  102. }\
  103. } while (0)
  104. //static Logger gLogger;
  105. //构建Logger
  106. class Logger : public ILogger
  107. {
  108. void log(Severity severity, const char* msg) noexcept override
  109. {
  110. // suppress info-level messages
  111. if (severity <= Severity::kWARNING)
  112. std::cout << msg << std::endl;
  113. }
  114. } gLogger;
  115. // Creat the engine using only the API and not any parser.
  116. ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config)
  117. {
  118. INetworkDefinition* network = builder->createNetworkV2(1U); //此处重点1U为OU就有问题
  119. IParser* parser = createParser(*network, gLogger);
  120. parser->parseFromFile(onnx_path.c_str(), static_cast<int32_t>(ILogger::Severity::kWARNING));
  121. //解析有错误将返回
  122. for (int32_t i = 0; i < parser->getNbErrors(); ++i) { std::cout << parser->getError(i)->desc() << std::endl; }
  123. std::cout << "successfully parse the onnx model" << std::endl;
  124. // Build engine
  125. builder->setMaxBatchSize(maxBatchSize);
  126. config->setMaxWorkspaceSize(1 << 20);
  127. //config->setFlag(nvinfer1::BuilderFlag::kFP16); // 设置精度计算
  128. //config->setFlag(nvinfer1::BuilderFlag::kINT8);
  129. ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
  130. std::cout << "successfully convert onnx to engine!!! " << std::endl;
  131. //销毁
  132. network->destroy();
  133. //parser->destroy();
  134. return engine;
  135. }
  136. void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream)
  137. {
  138. // Create builder
  139. IBuilder* builder = createInferBuilder(gLogger);
  140. IBuilderConfig* config = builder->createBuilderConfig();
  141. // Create model to populate the network, then set the outputs and create an engine
  142. ICudaEngine* engine = createEngine(maxBatchSize, builder, config);
  143. assert(engine != nullptr);
  144. // Serialize the engine
  145. (*modelStream) = engine->serialize();
  146. // Close everything down
  147. engine->destroy();
  148. builder->destroy();
  149. config->destroy();
  150. }
  151. int get_trtengine() {
  152. IHostMemory* modelStream{ nullptr };
  153. APIToModel(1, &modelStream);
  154. assert(modelStream != nullptr);
  155. std::ofstream p(engine_path, std::ios::binary);
  156. if (!p)
  157. {
  158. std::cerr << "could not open plan output file" << std::endl;
  159. return -1;
  160. }
  161. p.write(reinterpret_cast<const char*>(modelStream->data()), modelStream->size());
  162. modelStream->destroy();
  163. return 0;
  164. }
  165. void doInference(IExecutionContext& context, float* input, float* output, int batchSize)
  166. {
  167. const ICudaEngine& engine = context.getEngine();
  168. // Pointers to input and output device buffers to pass to engine.
  169. // Engine requires exactly IEngine::getNbBindings() number of buffers.
  170. assert(engine.getNbBindings() == 2);
  171. void* buffers[2];
  172. // In order to bind the buffers, we need to know the names of the input and output tensors.
  173. // Note that indices are guaranteed to be less than IEngine::getNbBindings()
  174. const int inputIndex = engine.getBindingIndex(INPUT_BLOB_NAME);
  175. const int outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME);
  176. //std::cout<<inputIndex<<" "<<outputIndex<<std::endl;
  177. //const int inputIndex = 0;
  178. //const int outputIndex = 1;
  179. // Create GPU buffers on device
  180. cudaMalloc(&buffers[inputIndex], batchSize * 3 * INPUT_H * INPUT_W * sizeof(float));
  181. cudaMalloc(&buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float));
  182. // Create stream
  183. cudaStream_t stream;
  184. CHECK(cudaStreamCreate(&stream));
  185. // DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host
  186. CHECK(cudaMemcpyAsync(buffers[inputIndex], input, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream));
  187. context.enqueue(batchSize, buffers, stream, nullptr);
  188. //std::cout<<buffers[outputIndex+1]<<std::endl;
  189. CHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream));
  190. cudaStreamSynchronize(stream);
  191. // Release stream and buffers
  192. cudaStreamDestroy(stream);
  193. CHECK(cudaFree(buffers[inputIndex]));
  194. CHECK(cudaFree(buffers[outputIndex]));
  195. }
  196. //加工图片变成拥有batch的输入, tensorrt输入需要的格式,为一个维度
  197. void ProcessImage(cv::Mat image, float input_data[]) {
  198. //只处理一张图片,总之结果为一维[batch*3*INPUT_W*INPUT_H]
  199. //以下代码为投机取巧了
  200. cv::Mat resize_img ;
  201. cv::resize(image, resize_img, cv::Size(INPUT_W, INPUT_H), 0, 0, cv::INTER_LINEAR);
  202. std::vector<cv::Mat> InputImage;
  203. InputImage.push_back(resize_img);
  204. int ImgCount = InputImage.size();
  205. //float input_data[BatchSize * 3 * INPUT_H * INPUT_W];
  206. for (int b = 0; b < ImgCount; b++) {
  207. cv::Mat img = InputImage.at(b);
  208. int w = img.cols;
  209. int h = img.rows;
  210. int i = 0;
  211. for (int row = 0; row < h; ++row) {
  212. uchar* uc_pixel = img.data + row * img.step;
  213. for (int col = 0; col < INPUT_W; ++col) {
  214. input_data[b * 3 * INPUT_H * INPUT_W + i] = (float)uc_pixel[2] / 255.0;
  215. input_data[b * 3 * INPUT_H * INPUT_W + i + INPUT_H * INPUT_W] = (float)uc_pixel[1] / 255.0;
  216. input_data[b * 3 * INPUT_H * INPUT_W + i + 2 * INPUT_H * INPUT_W] = (float)uc_pixel[0] / 255.0;
  217. uc_pixel += 3;
  218. ++i;
  219. }
  220. }
  221. }
  222. }
  223. //********************************************** NMS code **********************************//
  224. float iou(Bbox box1, Bbox box2) {
  225. int x1 = max(box1.x, box2.x);
  226. int y1 = max(box1.y, box2.y);
  227. int x2 = min(box1.x + box1.w, box2.x + box2.w);
  228. int y2 = min(box1.y + box1.h, box2.y + box2.h);
  229. int w = max(0, x2 - x1);
  230. int h = max(0, y2 - y1);
  231. float over_area = w * h;
  232. return over_area / (box1.w * box1.h + box2.w * box2.h - over_area);
  233. }
  234. int get_max_index(vector<Detection> pre_detection) {
  235. //获得最佳置信度的值,并返回对应的索引值
  236. int index;
  237. float conf;
  238. if (pre_detection.size() > 0) {
  239. index = 0;
  240. conf = pre_detection.at(0).conf;
  241. for (int i = 0; i < pre_detection.size(); i++) {
  242. if (conf < pre_detection.at(i).conf) {
  243. index = i;
  244. conf = pre_detection.at(i).conf;
  245. }
  246. }
  247. return index;
  248. }
  249. else {
  250. return -1;
  251. }
  252. }
  253. bool judge_in_lst(int index, vector<int> index_lst) {
  254. //若index在列表index_lst中则返回true,否则返回false
  255. if (index_lst.size() > 0) {
  256. for (int i = 0; i < index_lst.size(); i++) {
  257. if (index == index_lst.at(i)) {
  258. return true;
  259. }
  260. }
  261. }
  262. return false;
  263. }
  264. vector<int> nms(vector<Detection> pre_detection, float iou_thr)
  265. {
  266. /*
  267. 返回需保存box的pre_detection对应位置索引值
  268. */
  269. int index;
  270. vector<Detection> pre_detection_new;
  271. //Detection det_best;
  272. Bbox box_best, box;
  273. float iou_value;
  274. vector<int> keep_index;
  275. vector<int> del_index;
  276. bool keep_bool;
  277. bool del_bool;
  278. int rr = 0;
  279. int zz = 0;
  280. if (pre_detection.size() > 0) {
  281. pre_detection_new.clear();
  282. // 循环将预测结果建立索引
  283. for (int i = 0; i < pre_detection.size(); i++) {
  284. pre_detection.at(i).index = i;
  285. pre_detection_new.push_back(pre_detection.at(i));
  286. }
  287. //循环遍历获得保留box位置索引-相对输入pre_detection位置
  288. while (pre_detection_new.size() > 0) {
  289. index = get_max_index(pre_detection_new);
  290. if (index >= 0) {
  291. keep_index.push_back(pre_detection_new.at(index).index); //保留索引位置
  292. // 更新最佳保留box
  293. box_best.x = pre_detection_new.at(index).bbox[0];
  294. box_best.y = pre_detection_new.at(index).bbox[1];
  295. box_best.w = pre_detection_new.at(index).bbox[2];
  296. box_best.h = pre_detection_new.at(index).bbox[3];
  297. for (int j = 0; j < pre_detection.size(); j++) {
  298. keep_bool = judge_in_lst(pre_detection.at(j).index, keep_index);
  299. del_bool = judge_in_lst(pre_detection.at(j).index, del_index);
  300. if ((!keep_bool) && (!del_bool)) { //不在keep_index与del_index才计算iou
  301. box.x = pre_detection.at(j).bbox[0];
  302. box.y = pre_detection.at(j).bbox[1];
  303. box.w = pre_detection.at(j).bbox[2];
  304. box.h = pre_detection.at(j).bbox[3];
  305. iou_value = iou(box_best, box);
  306. if (iou_value > iou_thr) {
  307. del_index.push_back(j); //记录大于阈值将删除对应的位置
  308. }
  309. }
  310. }
  311. //更新pre_detection_new
  312. pre_detection_new.clear();
  313. for (int j = 0; j < pre_detection.size(); j++) {
  314. keep_bool = judge_in_lst(pre_detection.at(j).index, keep_index);
  315. del_bool = judge_in_lst(pre_detection.at(j).index, del_index);
  316. if ((!keep_bool) && (!del_bool)) {
  317. pre_detection_new.push_back(pre_detection.at(j));
  318. }
  319. }
  320. }
  321. }
  322. }
  323. del_index.clear();
  324. del_index.shrink_to_fit();
  325. pre_detection_new.clear();
  326. pre_detection_new.shrink_to_fit();
  327. return keep_index;
  328. }
  329. void postprocess(float* prob,vector<Detection> &results,float conf_thr = 0.2, float nms_thr = 0.4) {
  330. /*
  331. #####################此函数处理一张图预测结果#########################
  332. prob为[x y w h score multi-pre] 如80类-->(1,anchor_num,85)
  333. */
  334. vector<Detection> pre_results;
  335. vector<int> nms_keep_index;
  336. bool keep_bool;
  337. Detection pre_res;
  338. float conf;
  339. int tmp_idx;
  340. float tmp_cls_score;
  341. for (int i = 0; i < anchor_output_num; i++) {
  342. tmp_idx = i * (cls_num + 5);
  343. pre_res.bbox[0] = prob[tmp_idx + 0];
  344. pre_res.bbox[1] = prob[tmp_idx + 1];
  345. pre_res.bbox[2] = prob[tmp_idx + 2];
  346. pre_res.bbox[3] = prob[tmp_idx + 3];
  347. conf = prob[tmp_idx + 4]; //是为目标的置信度
  348. tmp_cls_score = prob[tmp_idx + 5] * conf;
  349. pre_res.class_id = 0;
  350. pre_res.conf = tmp_cls_score;
  351. for (int j = 1; j < cls_num; j++) {
  352. tmp_idx = i * (cls_num + 5) + 5 + j; //获得对应类别索引
  353. if (tmp_cls_score < prob[tmp_idx] * conf)
  354. {
  355. tmp_cls_score = prob[tmp_idx] * conf;
  356. pre_res.class_id = j;
  357. pre_res.conf = tmp_cls_score;
  358. }
  359. }
  360. if (conf >= conf_thr) {
  361. pre_results.push_back(pre_res);
  362. }
  363. }
  364. //使用nms
  365. nms_keep_index=nms(pre_results,nms_thr);
  366. for (int i = 0; i < pre_results.size(); i++) {
  367. keep_bool = judge_in_lst(i, nms_keep_index);
  368. if (keep_bool) {
  369. results.push_back(pre_results.at(i));
  370. }
  371. }
  372. pre_results.clear();
  373. pre_results.shrink_to_fit();
  374. nms_keep_index.clear();
  375. nms_keep_index.shrink_to_fit();
  376. }
  377. cv::Mat draw_rect(cv::Mat image, vector<Detection> results) {
  378. /*
  379. image 为图像
  380. struct Detection {
  381. float bbox[4]; //center_x center_y w h
  382. float conf; // 置信度
  383. int class_id; //类别id
  384. int index; //可忽略
  385. };
  386. */
  387. float x;
  388. float y;
  389. float w;
  390. float h;
  391. string info;
  392. cv::Rect rect;
  393. for (int i = 0; i < results.size(); i++) {
  394. x = results.at(i).bbox[0];
  395. y= results.at(i).bbox[1];
  396. w= results.at(i).bbox[2];
  397. h=results.at(i).bbox[3];
  398. x = (int)(x - w / 2);
  399. y = (int)(y - h / 2);
  400. w = (int)w;
  401. h = (int)h;
  402. info = "id:";
  403. info.append(to_string(results.at(i).class_id));
  404. info.append(" s:");
  405. info.append(to_string((int)(results.at(i).conf*100) ) );
  406. info.append("%");
  407. rect= cv::Rect(x, y, w, h);
  408. if(results.at(i).class_id == 0){ // red light
  409. cv::rectangle(image, rect, cv::Scalar(0, 0, 255), 2, 1, 0);//矩形的两个顶点,两个顶点都包括在矩形内部
  410. cv::putText(image, info, cv::Point(x, y), cv::FONT_HERSHEY_SIMPLEX, 0.4, cv::Scalar(0, 0, 255), 0.6, 1, false);
  411. }else if(results.at(i).class_id == 1){ // green light
  412. cv::rectangle(image, rect, cv::Scalar(0, 255, 0), 2, 1, 0);//矩形的两个顶点,两个顶点都包括在矩形内部
  413. cv::putText(image, info, cv::Point(x, y), cv::FONT_HERSHEY_SIMPLEX, 0.4, cv::Scalar(0, 255, 0), 0.6, 1, false);
  414. }else if(results.at(i).class_id == 2){ // yellow light
  415. cv::rectangle(image, rect, cv::Scalar(0, 255, 255), 2, 1, 0);//矩形的两个顶点,两个顶点都包括在矩形内部
  416. cv::putText(image, info, cv::Point(x, y), cv::FONT_HERSHEY_SIMPLEX, 0.4, cv::Scalar(0, 255, 255), 0.6, 1, false);
  417. }else{
  418. cv::rectangle(image, rect, cv::Scalar(255, 255, 255), 2, 1, 0);//矩形的两个顶点,两个顶点都包括在矩形内部
  419. cv::putText(image, info, cv::Point(x, y), cv::FONT_HERSHEY_SIMPLEX, 0.4, cv::Scalar(255, 255, 255), 0.6, 1, false);
  420. }
  421. }
  422. return image;
  423. }
  424. bool LoadEngine(const std::string engine_path){
  425. //加载engine引擎
  426. char* trtModelStream{ nullptr };
  427. size_t size{ 0 };
  428. std::ifstream file(engine_path, std::ios::binary);
  429. if(!file)
  430. {
  431. cout<<engine_path<<" not found!"<<endl;
  432. return false;
  433. }
  434. if (file.good()) {
  435. file.seekg(0, file.end);
  436. size = file.tellg();
  437. file.seekg(0, file.beg);
  438. trtModelStream = new char[size];
  439. assert(trtModelStream);
  440. file.read(trtModelStream, size);
  441. file.close();
  442. }
  443. //反序列为engine,创建context
  444. runtime = createInferRuntime(gLogger);
  445. assert(runtime != nullptr);
  446. #ifdef UBUNTU_22_04
  447. engine = runtime->deserializeCudaEngine(trtModelStream, size);
  448. #else
  449. engine = runtime->deserializeCudaEngine(trtModelStream, size, nullptr);
  450. #endif
  451. //assert(engine != nullptr);
  452. if(engine == nullptr)
  453. return false;
  454. context = engine->createExecutionContext();
  455. assert(context != nullptr);
  456. delete[] trtModelStream;
  457. //在主机上分配页锁定内存
  458. CHECK(cudaHostAlloc((void **)&prob, OUTPUT_SIZE * sizeof(float), cudaHostAllocDefault));
  459. return true;
  460. }
  461. void shareLightMsg(vector<Detection> results)
  462. {
  463. iv::vision::Lightarray light_array; //向共享内存传结果
  464. for (int i = 0; i < results.size(); i++)
  465. {
  466. float x = results.at(i).bbox[0];
  467. float y = results.at(i).bbox[1];
  468. float w = results.at(i).bbox[2];
  469. float h = results.at(i).bbox[3];
  470. /*---------------protobuf----------------*/
  471. iv::vision::Light *light = light_array.add_light();
  472. iv::vision::Center light_center;
  473. light_center.set_x(x);
  474. light_center.set_y(y);
  475. light->mutable_center()->CopyFrom(light_center);
  476. light->set_index(i+1);
  477. int light_type = 255;
  478. if (results.at(i).class_id == 0) //红色
  479. light_type = 2;
  480. else if(results.at(i).class_id == 1) //绿色
  481. light_type = 1;
  482. else if(results.at(i).class_id == 2) //黄色
  483. light_type = 3;
  484. light->set_type(light_type);
  485. cout<<"light color: "<<light_type<<endl;
  486. if(ivlog_flag)
  487. givlog->verbose("light color: %d",light_type);
  488. }
  489. int size = light_array.ByteSize();
  490. char * strdata = new char[light_array.ByteSize()];
  491. if(light_array.SerializeToArray(strdata, size))
  492. {
  493. iv::modulecomm::ModuleSendMsg(gpdetect, strdata, size);
  494. }
  495. else
  496. {
  497. std::cout<<"light_array serialize error."<<std::endl;
  498. }
  499. light_array.Clear();
  500. delete strdata;
  501. /*--------------------test ParseFromArray-------------------*/
  502. // iv::vision::Lightarray light_array1;
  503. // light_array1.ParseFromArray(strdata,size);
  504. // cout<<"parsefromarray:"<<std::endl;
  505. // cout<<"light_size:"<<light_array1.light_size()<<endl;
  506. // for (int i=0;i<light_array1.light_size();i++) {
  507. // std::cout<<"index:"<<light_array1.light(i).index()<<" type:"
  508. // <<light_array1.light(i).type()
  509. // <<std::endl;
  510. // }
  511. }
  512. void infer(cv::Mat img,vector<Detection> &results) {
  513. // 处理图片为固定输出
  514. auto start = std::chrono::system_clock::now(); //时间函数
  515. static float data[3 * INPUT_H * INPUT_W];
  516. ProcessImage(img, data);
  517. auto end = std::chrono::system_clock::now();
  518. //time_read_img = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() + time_read_img;
  519. //cout<<"read img time: "<<std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count()<<"ms"<<endl;
  520. //Run inference
  521. start = std::chrono::system_clock::now(); //时间函数
  522. //cout<<"doinference"<<endl;
  523. doInference(*context, data, prob, 1);
  524. end = std::chrono::system_clock::now();
  525. //time_infer = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() + time_infer;
  526. std::cout <<"doinference: "<< std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
  527. postprocess(prob, results, conf_thr, nms_thr);
  528. //cout << "ok" << endl;
  529. //time_num++;
  530. }
  531. void exitfunc()
  532. {
  533. return;
  534. }
  535. //读取视频数据
  536. void ReadFunc(int n)
  537. {
  538. cv::VideoCapture cap(video_path);
  539. if(!cap.isOpened())
  540. {
  541. cout<<"camera failed to open"<<endl;
  542. }
  543. while(1)
  544. {
  545. cv::Mat frame;
  546. //读视频的时候加上,读摄像头去掉
  547. if(imageBuffer->isFull())
  548. {
  549. continue;
  550. }
  551. if(cap.read(frame))
  552. {
  553. if(calibrationstate)
  554. cv::remap(frame,frame,map1,map2,cv::INTER_LINEAR,cv::BORDER_CONSTANT);
  555. if(cropstate)
  556. frame = frame(crop_height,crop_width);
  557. imageBuffer->add(frame);
  558. }
  559. else
  560. {
  561. std::this_thread::sleep_for(std::chrono::milliseconds(1));
  562. }
  563. }
  564. }
  565. void Listenpic(const char * strdata,const unsigned int nSize,const unsigned int index,const QDateTime * dt,const char * strmemname)
  566. {
  567. if(nSize<1000)return;
  568. iv::vision::rawpic pic;
  569. if(false == pic.ParseFromArray(strdata,nSize))
  570. {
  571. std::cout<<"picview Listenpic fail."<<std::endl;
  572. return;
  573. }
  574. cv::Mat mat(pic.height(),pic.width(),pic.mattype());
  575. if(pic.type() == 1)
  576. memcpy(mat.data,pic.picdata().data(),mat.rows*mat.cols*mat.elemSize());
  577. else
  578. {
  579. // mat.release();
  580. std::vector<unsigned char> buff(pic.picdata().data(),pic.picdata().data() + pic.picdata().size());
  581. mat = cv::imdecode(buff,cv::IMREAD_COLOR);
  582. }
  583. if(calibrationstate)
  584. cv::remap(mat,mat,map1,map2,cv::INTER_LINEAR,cv::BORDER_CONSTANT);
  585. if(cropstate)
  586. mat = mat(crop_height,crop_width);
  587. imageBuffer->add(mat);
  588. mat.release();
  589. }
  590. //从共享内存中获取signal
  591. void Listensignal(const char * strdata,const unsigned int nSize,const unsigned int index,const QDateTime * dt,const char * strmemname)
  592. {
  593. if(nSize<1)return;
  594. iv::brain::brainstate brain_state;
  595. if(false == brain_state.ParseFromArray(strdata,nSize))
  596. {
  597. std::cout<<"Listen signal fail."<<std::endl;
  598. return;
  599. }
  600. cout<<"brain_state traficlightstart: "<<brain_state.mbtraficlightstart()<<endl;
  601. lightstart = brain_state.mbtraficlightstart();
  602. }
  603. bool comp(const od::TrackingBox &a, const od::TrackingBox &b) {
  604. return a.id < b.id;
  605. }
  606. int main(int argc, char** argv)
  607. {
  608. showversion("yolov5s");
  609. QCoreApplication a(argc, argv);
  610. gfault = new iv::Ivfault("traffic_light_detection");
  611. givlog = new iv::Ivlog("traffic_light_detection");
  612. gfault->SetFaultState(0,0,"yolov5 initialize.");
  613. // if(argc==3)
  614. // {
  615. // test_video = (strcmp(argv[1], "true") == 0)?true:false;
  616. // video_path = argv[2];
  617. // }
  618. // if(argc==2)
  619. // {
  620. // test_video = (strcmp(argv[1], "true") == 0)?true:false;
  621. // }
  622. cv::FileStorage config(config_file, cv::FileStorage::READ);
  623. bool config_isOpened = config.isOpened();
  624. //const char* onnx_path_;
  625. if(config_isOpened)
  626. {
  627. onnx_path = string(config["model"]["onnx_path"]);
  628. engine_path = string(config["model"]["engine_path"]);
  629. imshow_flag = (string(config["imshow_flag"]) == "true");
  630. ivlog_flag = (string(config["ivlog_flag"]) == "true");
  631. test_video = (string(config["test_video"]["open"]) == "true");
  632. video_path = string(config["test_video"]["video_path"]);
  633. lightstart = (string(config["lightstart"]) == "true");
  634. trackstate = (string(config["trackstate"]) == "true");
  635. conf_thr = float(config["conf_thr"]);
  636. nms_thr = float(config["nms_thr"]);
  637. calibrationstate = (string(config["camera"]["calibrationstate"]) == "true");
  638. calibration_yamlpath = string(config["camera"]["calibration_yamlpath"]);
  639. cropstate = (string(config["camera"]["cropstate"]) == "true");
  640. crop_height = cv::Range(int(config["camera"]["crop_height"]["start"]),
  641. int(config["camera"]["crop_height"]["end"]));
  642. crop_width = cv::Range(int(config["camera"]["crop_width"]["start"]),
  643. int(config["camera"]["crop_width"]["end"]));
  644. }
  645. config.release();
  646. if(test_video)
  647. std::thread * readthread = new std::thread(ReadFunc,1);
  648. else
  649. gpcamera= iv::modulecomm::RegisterRecv(&cameraname[0],Listenpic);
  650. //================================== camera calib init ==========================
  651. if (calibrationstate)
  652. {
  653. cv::FileStorage calib_file(calibration_yamlpath, cv::FileStorage::READ);
  654. calib_file["cameraMatrix"]>>camera_matrix;
  655. calib_file["distCoeffs"]>>dist_coe;
  656. cv::Mat R = cv::Mat::eye(3, 3, CV_64F);
  657. cv::Size imgsize=cv::Size(1280,720);
  658. cv::initUndistortRectifyMap(camera_matrix, dist_coe, R, camera_matrix,imgsize,CV_16SC2,map1,map2);
  659. }
  660. //==============================================================================
  661. if(!LoadEngine(engine_path))
  662. {
  663. cout<<"Build engine to "<< engine_path <<endl;
  664. get_trtengine();
  665. cout << "Build engine done!"<<endl;
  666. cout<<"Reload engine from "<< engine_path <<endl;
  667. LoadEngine(engine_path);
  668. }
  669. gpbrain = iv::modulecomm::RegisterRecv(&brainname[0],Listensignal);
  670. gpdetect = iv::modulecomm::RegisterSend(&detectname[0],10000,1);
  671. while(1)
  672. {
  673. if(lightstart)
  674. {
  675. std::cout<<"------start program------"<<std::endl;
  676. vector<KalmanTracker> trackers_90;
  677. KalmanTracker::kf_count = 0; // tracking id relies on this, so we have to reset it in each seq.
  678. int frame_count = 0;
  679. double waittime = (double)cv::getTickCount();
  680. while (1)
  681. {
  682. if(imageBuffer->isEmpty())
  683. {
  684. double waittotal = (double)cv::getTickCount() - waittime;
  685. double totaltime = waittotal/cv::getTickFrequency();
  686. // if(totaltime>10.0)
  687. // {
  688. // cout<<"Cant't get frame and quit"<<endl;
  689. // lightstart = false;
  690. // cv::destroyAllWindows();
  691. // std::cout<<"------end program------"<<std::endl;
  692. // break;
  693. // }
  694. cout<<"Wait for frame "<<totaltime<<"s"<<endl;
  695. if(ivlog_flag)
  696. givlog->verbose("Wait for frame %f s",totaltime);
  697. std::this_thread::sleep_for(std::chrono::milliseconds(10));
  698. continue;
  699. }
  700. auto start = std::chrono::system_clock::now(); //时间函数
  701. cv::Mat frame;
  702. imageBuffer->consume(frame);
  703. frame_count++;
  704. cv::Mat res_img; // result image
  705. vector<Detection> results;
  706. vector<Detection>results_track;
  707. //================================== infer ==========================
  708. infer(frame,results);
  709. //================================== track ==========================
  710. if (trackstate)
  711. {
  712. auto start_track = std::chrono::system_clock::now(); //时间函数
  713. od::bbox_t bbox_t_90; //转成跟踪格式
  714. vector<od::bbox_t> outs_90;
  715. for (int i = 0; i < results.size(); i++)
  716. {
  717. //-------------判断红绿灯是否为横向,width=(x1-x2),height=(y1-y2)-----------
  718. bbox_t_90.x = results.at(i).bbox[0];
  719. bbox_t_90.y = results.at(i).bbox[1];
  720. bbox_t_90.w = results.at(i).bbox[2];
  721. bbox_t_90.h = results.at(i).bbox[3];
  722. bbox_t_90.prob = results.at(i).conf;
  723. bbox_t_90.obj_id = results.at(i).class_id;
  724. outs_90.push_back(bbox_t_90);
  725. }
  726. vector<od::TrackingBox>track_result_90;
  727. bool track_flag_90 = od::TrackObstacle(frame_count,trackers_90,outs_90,track_result_90);
  728. //sort(track_result_90.begin(), track_result_90.end(), comp); //track id 本来就是由大到小
  729. for(unsigned int i=0;i < track_result_90.size(); i++)
  730. {
  731. Detection obstacle;
  732. obstacle.bbox[0] = track_result_90[i].box.x;
  733. obstacle.bbox[1] = track_result_90[i].box.y;
  734. obstacle.bbox[2] = track_result_90[i].box.width;
  735. obstacle.bbox[3] = track_result_90[i].box.height;
  736. //cout<<"11111: "<<track_result_90[i].id<<endl;
  737. //通过判断5帧数输出颜色
  738. vector<int> class_history;
  739. class_history = track_result_90[i].class_history;
  740. if(class_history.size()>0)
  741. {
  742. vector<int> color_num(3);
  743. for(int j=0;j<class_history.size();j++)
  744. {
  745. int class_id = class_history[j];
  746. color_num[class_id] += 1;
  747. }
  748. std::vector<int>::iterator biggest = std::max_element(std::begin(color_num),std::end(color_num));
  749. int maxindex = std::distance(std::begin(color_num),biggest);
  750. obstacle.class_id = maxindex;
  751. }
  752. else {obstacle.class_id = track_result_90[i].class_id;}
  753. obstacle.conf = track_result_90[i].prob;
  754. results_track.push_back(obstacle);
  755. }
  756. auto end_track = std::chrono::system_clock::now(); //时间函数
  757. //std::cout <<"track: "<< std::chrono::duration_cast<std::chrono::milliseconds>(end_track - start_track).count() << "ms" << std::endl;
  758. }
  759. //================================== track ==========================
  760. vector<Detection>results_final;
  761. results_final = (trackstate)?results_track:results;
  762. shareLightMsg(results_final);
  763. auto end = std::chrono::system_clock::now(); //时间函数
  764. std::cout <<"total time traffic: "<< std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
  765. //显示检测结果,检测结果没有返回到原始图像大小,所以需要resize到网络输入尺寸
  766. cv::resize(frame, frame, cv::Size(INPUT_W, INPUT_H), 0, 0, cv::INTER_LINEAR);
  767. res_img=draw_rect(frame,results_final);
  768. std::cout<<"light number: "<<results_final.size()<<std::endl;
  769. if(ivlog_flag)
  770. givlog->verbose("light number: %d",results_final.size());
  771. results.clear();
  772. results_track.clear();
  773. results_final.clear();
  774. if (imshow_flag)
  775. {
  776. cv::namedWindow("Result",cv::WINDOW_NORMAL);
  777. cv::imshow("Result",frame);
  778. if(cv::waitKey(10) == 'q')
  779. {
  780. cv::destroyAllWindows();
  781. //yolo_context->destroy();
  782. //start_up = false;
  783. break;
  784. }
  785. if(cv::waitKey(1) == 's')
  786. cv::waitKey(0);
  787. }
  788. //writer << frame;
  789. waittime = (double)cv::getTickCount();
  790. if(!lightstart)
  791. {
  792. cv::destroyAllWindows();
  793. //yolo_context->destroy();
  794. //start_up = false;
  795. std::cout<<"------end program------"<<std::endl;
  796. break;
  797. }
  798. }
  799. }
  800. else
  801. {
  802. cout<<"Wait for lightstart==true"<<endl;
  803. if(ivlog_flag)
  804. givlog->verbose("Wait for lightstart==true");
  805. std::this_thread::sleep_for(std::chrono::milliseconds(10));
  806. }
  807. }
  808. // Destroy the engine
  809. context->destroy();
  810. engine->destroy();
  811. runtime->destroy();
  812. return 0;
  813. }