segmenter.cpp 3.0 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586
  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 <fstream>
  16. #include <iostream>
  17. #include <string>
  18. #include <vector>
  19. #include "include/paddlex/paddlex.h"
  20. #include "include/paddlex/visualize.h"
  21. DEFINE_string(model_dir, "", "Path of inference model");
  22. DEFINE_bool(use_gpu, false, "Infering with GPU or CPU");
  23. DEFINE_bool(use_trt, false, "Infering with TensorRT");
  24. DEFINE_int32(gpu_id, 0, "GPU card id");
  25. DEFINE_string(image, "", "Path of test image file");
  26. DEFINE_string(image_list, "", "Path of test image list file");
  27. DEFINE_string(save_dir, "output", "Path to save visualized image");
  28. int main(int argc, char** argv) {
  29. // 解析命令行参数
  30. google::ParseCommandLineFlags(&argc, &argv, true);
  31. if (FLAGS_model_dir == "") {
  32. std::cerr << "--model_dir need to be defined" << std::endl;
  33. return -1;
  34. }
  35. if (FLAGS_image == "" & FLAGS_image_list == "") {
  36. std::cerr << "--image or --image_list need to be defined" << std::endl;
  37. return -1;
  38. }
  39. // 加载模型
  40. PaddleX::Model model;
  41. model.Init(FLAGS_model_dir, FLAGS_use_gpu, FLAGS_use_trt, FLAGS_gpu_id);
  42. auto colormap = PaddleX::GenerateColorMap(model.labels.size());
  43. // 进行预测
  44. if (FLAGS_image_list != "") {
  45. std::ifstream inf(FLAGS_image_list);
  46. if (!inf) {
  47. std::cerr << "Fail to open file " << FLAGS_image_list << std::endl;
  48. return -1;
  49. }
  50. std::string image_path;
  51. while (getline(inf, image_path)) {
  52. PaddleX::SegResult result;
  53. cv::Mat im = cv::imread(image_path, 1);
  54. model.predict(im, &result);
  55. // 可视化
  56. cv::Mat vis_img =
  57. PaddleX::VisualizeSeg(im, result, model.labels, colormap);
  58. std::string save_path =
  59. PaddleX::generate_save_path(FLAGS_save_dir, image_path);
  60. cv::imwrite(save_path, vis_img);
  61. result.clear();
  62. std::cout << "Visualized output saved as " << save_path << std::endl;
  63. }
  64. } else {
  65. PaddleX::SegResult result;
  66. cv::Mat im = cv::imread(FLAGS_image, 1);
  67. model.predict(im, &result);
  68. // 可视化
  69. cv::Mat vis_img = PaddleX::VisualizeSeg(im, result, model.labels, colormap);
  70. std::string save_path =
  71. PaddleX::generate_save_path(FLAGS_save_dir, FLAGS_image);
  72. cv::imwrite(save_path, vis_img);
  73. result.clear();
  74. std::cout << "Visualized output saved as " << save_path << std::endl;
  75. }
  76. return 0;
  77. }