meter_reader.cpp 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326
  1. // Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License");
  4. // you may not use this file except in compliance with the License.
  5. // You may obtain a copy of the License at
  6. //
  7. // http://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS,
  11. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. // See the License for the specific language governing permissions and
  13. // limitations under the License.
  14. #include <glog/logging.h>
  15. #include <omp.h>
  16. #include <algorithm>
  17. #include <chrono> // NOLINT
  18. #include <iostream>
  19. #include <vector>
  20. #include <utility>
  21. #include <limits>
  22. #include <opencv2/opencv.hpp>
  23. #include <opencv2/highgui.hpp>
  24. #include <opencv2/core/core.hpp>
  25. #include "meter_reader/global.h"
  26. #include "meter_reader/postprocess.h"
  27. #include "include/paddlex/paddlex.h"
  28. #include "include/paddlex/visualize.h"
  29. using namespace std::chrono; // NOLINT
  30. DEFINE_string(det_model_dir, "", "Path of detection inference model");
  31. DEFINE_string(seg_model_dir, "", "Path of segmentation inference model");
  32. DEFINE_bool(use_gpu, false, "Infering with GPU or CPU");
  33. DEFINE_bool(use_trt, false, "Infering with TensorRT");
  34. DEFINE_bool(use_camera, false, "Infering with Camera");
  35. DEFINE_bool(use_erode, true, "Eroding predicted label map");
  36. DEFINE_int32(gpu_id, 0, "GPU card id");
  37. DEFINE_int32(camera_id, 0, "Camera id");
  38. DEFINE_int32(thread_num,
  39. omp_get_num_procs(),
  40. "Number of preprocessing threads");
  41. DEFINE_int32(erode_kernel, true, "Eroding kernel size");
  42. DEFINE_int32(seg_batch_size, 2, "Batch size of segmentation infering");
  43. DEFINE_string(det_key, "", "Detector key of encryption");
  44. DEFINE_string(seg_key, "", "Segmenter model key of encryption");
  45. DEFINE_string(image, "", "Path of test image file");
  46. DEFINE_string(image_list, "", "Path of test image list file");
  47. DEFINE_string(save_dir, "output", "Path to save visualized image");
  48. DEFINE_double(score_threshold, 0.5,
  49. "Detected bbox whose score is lower than this threshlod is filtered");
  50. void predict(const cv::Mat &input_image, PaddleX::Model *det_model,
  51. PaddleX::Model *seg_model, const std::string save_dir,
  52. const std::string image_path, const bool use_erode,
  53. const int erode_kernel, const int thread_num,
  54. const int seg_batch_size, const double threshold) {
  55. // Get detection results
  56. PaddleX::DetResult det_result;
  57. det_model->predict(input_image, &det_result);
  58. // Filter bbox whose score is lower than score_threshold
  59. PaddleX::DetResult filter_result;
  60. int num_bboxes = det_result.boxes.size();
  61. for (int i = 0; i < num_bboxes; ++i) {
  62. double score = det_result.boxes[i].score;
  63. if (score > threshold || score == threshold) {
  64. PaddleX::Box box;
  65. box.category_id = det_result.boxes[i].category_id;
  66. box.category = det_result.boxes[i].category;
  67. box.score = det_result.boxes[i].score;
  68. box.coordinate = det_result.boxes[i].coordinate;
  69. filter_result.boxes.push_back(std::move(box));
  70. }
  71. }
  72. int meter_num = filter_result.boxes.size();
  73. if (!meter_num) {
  74. std::cout << "Don't find any meter." << std::endl;
  75. return;
  76. }
  77. std::vector<std::vector<int64_t>> seg_result(meter_num);
  78. for (int i = 0; i < meter_num; i += seg_batch_size) {
  79. int im_vec_size =
  80. std::min(static_cast<int>(meter_num), i + seg_batch_size);
  81. std::vector<cv::Mat> meters_image(im_vec_size - i);
  82. int batch_thread_num = std::min(thread_num, im_vec_size - i);
  83. #pragma omp parallel for num_threads(batch_thread_num)
  84. for (int j = i; j < im_vec_size; ++j) {
  85. // Crop the bbox area
  86. int left = static_cast<int>(filter_result.boxes[j].coordinate[0]);
  87. int top = static_cast<int>(filter_result.boxes[j].coordinate[1]);
  88. int width = static_cast<int>(filter_result.boxes[j].coordinate[2]);
  89. int height = static_cast<int>(filter_result.boxes[j].coordinate[3]);
  90. int right = left + width - 1;
  91. int bottom = top + height - 1;
  92. cv::Mat sub_image = input_image(
  93. cv::Range(top, bottom + 1), cv::Range(left, right + 1));
  94. // Resize the image with shape (METER_SHAPE, METER_SHAPE)
  95. float scale_x =
  96. static_cast<float>(METER_SHAPE[0]) / static_cast<float>(sub_image.cols);
  97. float scale_y =
  98. static_cast<float>(METER_SHAPE[1]) / static_cast<float>(sub_image.rows);
  99. cv::resize(sub_image,
  100. sub_image,
  101. cv::Size(),
  102. scale_x,
  103. scale_y,
  104. cv::INTER_LINEAR);
  105. meters_image[j - i] = std::move(sub_image);
  106. }
  107. // Segment scales and point in each meter area
  108. std::vector<PaddleX::SegResult> batch_result(im_vec_size - i);
  109. seg_model->predict(meters_image, &batch_result, batch_thread_num);
  110. #pragma omp parallel for num_threads(batch_thread_num)
  111. for (int j = i; j < im_vec_size; ++j) {
  112. // Do image erosion for the predicted label map of each meter
  113. if (use_erode) {
  114. cv::Mat kernel(4, 4, CV_8U, cv::Scalar(1));
  115. std::vector<uint8_t> label_map(
  116. batch_result[j - i].label_map.data.begin(),
  117. batch_result[j - i].label_map.data.end());
  118. cv::Mat mask(batch_result[j - i].label_map.shape[0],
  119. batch_result[j - i].label_map.shape[1],
  120. CV_8UC1,
  121. label_map.data());
  122. cv::erode(mask, mask, kernel);
  123. std::vector<int64_t> map;
  124. if (mask.isContinuous()) {
  125. map.assign(mask.data, mask.data + mask.total() * mask.channels());
  126. } else {
  127. for (int r = 0; r < mask.rows; r++) {
  128. map.insert(map.end(),
  129. mask.ptr<int64_t>(r),
  130. mask.ptr<int64_t>(r) + mask.cols * mask.channels());
  131. }
  132. }
  133. seg_result[j] = std::move(map);
  134. } else {
  135. seg_result[j] = std::move(batch_result[j - i].label_map.data);
  136. }
  137. }
  138. }
  139. std::vector<READ_RESULT> read_results(meter_num);
  140. int all_thread_num = std::min(thread_num, meter_num);
  141. // The postprocess are done to get the point location relative to the scales
  142. read_process(seg_result, &read_results, all_thread_num);
  143. cv::Mat output_image = input_image.clone();
  144. for (int i = 0; i < meter_num; i++) {
  145. // Provide a digital readout according to point location relative
  146. // to the scales
  147. float result = 0;;
  148. if (read_results[i].scale_num > TYPE_THRESHOLD) {
  149. result = read_results[i].scales * meter_config[0].scale_value;
  150. } else {
  151. result = read_results[i].scales * meter_config[1].scale_value;
  152. }
  153. std::cout << "-- Meter " << i
  154. << " -- result: " << result
  155. << " --" << std::endl;
  156. // Visualize the results
  157. int lx = static_cast<int>(filter_result.boxes[i].coordinate[0]);
  158. int ly = static_cast<int>(filter_result.boxes[i].coordinate[1]);
  159. int w = static_cast<int>(filter_result.boxes[i].coordinate[2]);
  160. int h = static_cast<int>(filter_result.boxes[i].coordinate[3]);
  161. cv::Rect bounding_box = cv::Rect(lx, ly, w, h) &
  162. cv::Rect(0, 0, output_image.cols, output_image.rows);
  163. if (w > 0 && h > 0) {
  164. cv::Scalar color = cv::Scalar(237, 189, 101);
  165. cv::rectangle(output_image, bounding_box, color);
  166. cv::rectangle(output_image,
  167. cv::Point2d(lx, ly),
  168. cv::Point2d(lx + w, ly - 30),
  169. color, -1);
  170. std::string class_name = "Meter";
  171. cv::putText(output_image,
  172. class_name + " " + std::to_string(result),
  173. cv::Point2d(lx, ly-5),
  174. cv::FONT_HERSHEY_SIMPLEX,
  175. 1, cv::Scalar(255, 255, 255), 2);
  176. }
  177. }
  178. cv::Mat result_image;
  179. cv::Size resize_size(RESULT_SHAPE[0], RESULT_SHAPE[1]);
  180. cv::resize(output_image, result_image, resize_size, 0, 0, cv::INTER_LINEAR);
  181. std::string save_path = PaddleX::generate_save_path(save_dir, image_path);
  182. cv::imwrite(save_path, result_image);
  183. return;
  184. }
  185. int main(int argc, char **argv) {
  186. google::ParseCommandLineFlags(&argc, &argv, true);
  187. if (FLAGS_det_model_dir == "") {
  188. std::cerr << "--det_model_dir need to be defined" << std::endl;
  189. return -1;
  190. }
  191. if (FLAGS_seg_model_dir == "") {
  192. std::cerr << "--seg_model_dir need to be defined" << std::endl;
  193. return -1;
  194. }
  195. if (FLAGS_image == "" & FLAGS_image_list == "" & FLAGS_use_camera == false) {
  196. std::cerr << "--image or --image_list need to be defined "
  197. << "when the camera is not been used" << std::endl;
  198. return -1;
  199. }
  200. // Load model
  201. PaddleX::Model det_model;
  202. det_model.Init(FLAGS_det_model_dir, FLAGS_use_gpu, FLAGS_use_trt,
  203. FLAGS_gpu_id, FLAGS_det_key);
  204. PaddleX::Model seg_model;
  205. seg_model.Init(FLAGS_seg_model_dir, FLAGS_use_gpu, FLAGS_use_trt,
  206. FLAGS_gpu_id, FLAGS_seg_key);
  207. double total_running_time_s = 0.0;
  208. double total_imread_time_s = 0.0;
  209. int imgs = 1;
  210. if (FLAGS_use_camera) {
  211. cv::VideoCapture cap(FLAGS_camera_id);
  212. cap.set(CV_CAP_PROP_FRAME_WIDTH, IMAGE_SHAPE[0]);
  213. cap.set(CV_CAP_PROP_FRAME_HEIGHT, IMAGE_SHAPE[1]);
  214. if (!cap.isOpened()) {
  215. std::cout << "Open the camera unsuccessfully." << std::endl;
  216. return -1;
  217. }
  218. std::cout << "Open the camera successfully." << std::endl;
  219. while (1) {
  220. auto start = system_clock::now();
  221. cv::Mat im;
  222. cap >> im;
  223. auto imread_end = system_clock::now();
  224. std::cout << "-------------------------" << std::endl;
  225. std::cout << "Got a camera image." << std::endl;
  226. std::string ext_name = ".jpg";
  227. predict(im, &det_model, &seg_model, FLAGS_save_dir,
  228. std::to_string(imgs) + ext_name, FLAGS_use_erode,
  229. FLAGS_erode_kernel, FLAGS_thread_num,
  230. FLAGS_seg_batch_size, FLAGS_score_threshold);
  231. imgs++;
  232. auto imread_duration = duration_cast<microseconds>(imread_end - start);
  233. total_imread_time_s += static_cast<double>(imread_duration.count()) *
  234. microseconds::period::num /
  235. microseconds::period::den;
  236. auto end = system_clock::now();
  237. auto duration = duration_cast<microseconds>(end - start);
  238. total_running_time_s += static_cast<double>(duration.count()) *
  239. microseconds::period::num /
  240. microseconds::period::den;
  241. }
  242. cap.release();
  243. cv::destroyAllWindows();
  244. } else {
  245. if (FLAGS_image_list != "") {
  246. std::ifstream inf(FLAGS_image_list);
  247. if (!inf) {
  248. std::cerr << "Fail to open file " << FLAGS_image_list << std::endl;
  249. return -1;
  250. }
  251. std::string image_path;
  252. while (getline(inf, image_path)) {
  253. auto start = system_clock::now();
  254. cv::Mat im = cv::imread(image_path, 1);
  255. imgs++;
  256. auto imread_end = system_clock::now();
  257. predict(im, &det_model, &seg_model, FLAGS_save_dir,
  258. image_path, FLAGS_use_erode, FLAGS_erode_kernel,
  259. FLAGS_thread_num, FLAGS_seg_batch_size,
  260. FLAGS_score_threshold);
  261. auto imread_duration = duration_cast<microseconds>(imread_end - start);
  262. total_imread_time_s += static_cast<double>(imread_duration.count()) *
  263. microseconds::period::num /
  264. microseconds::period::den;
  265. auto end = system_clock::now();
  266. auto duration = duration_cast<microseconds>(end - start);
  267. total_running_time_s += static_cast<double>(duration.count()) *
  268. microseconds::period::num /
  269. microseconds::period::den;
  270. }
  271. } else {
  272. auto start = system_clock::now();
  273. cv::Mat im = cv::imread(FLAGS_image, 1);
  274. auto imread_end = system_clock::now();
  275. predict(im, &det_model, &seg_model, FLAGS_save_dir,
  276. FLAGS_image, FLAGS_use_erode, FLAGS_erode_kernel,
  277. FLAGS_thread_num, FLAGS_seg_batch_size,
  278. FLAGS_score_threshold);
  279. auto imread_duration = duration_cast<microseconds>(imread_end - start);
  280. total_imread_time_s += static_cast<double>(imread_duration.count()) *
  281. microseconds::period::num /
  282. microseconds::period::den;
  283. auto end = system_clock::now();
  284. auto duration = duration_cast<microseconds>(end - start);
  285. total_running_time_s += static_cast<double>(duration.count()) *
  286. microseconds::period::num /
  287. microseconds::period::den;
  288. }
  289. }
  290. std::cout << "Total running time: " << total_running_time_s
  291. << " s, average running time: " << total_running_time_s / imgs
  292. << " s/img, total read img time: " << total_imread_time_s
  293. << " s, average read time: " << total_imread_time_s / imgs
  294. << " s/img" << std::endl;
  295. return 0;
  296. }