preprocess_op.h 4.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169
  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. #pragma once
  15. #include <yaml-cpp/yaml.h>
  16. #include <vector>
  17. #include <string>
  18. #include <utility>
  19. #include <memory>
  20. #include <unordered_map>
  21. #include <opencv2/core/core.hpp>
  22. #include <opencv2/imgproc/imgproc.hpp>
  23. #include <opencv2/highgui/highgui.hpp>
  24. namespace PaddleDetection {
  25. // Object for storing all preprocessed data
  26. class ImageBlob {
  27. public:
  28. // Original image width and height
  29. std::vector<int> ori_im_size_;
  30. // Buffer for image data after preprocessing
  31. std::vector<float> im_data_;
  32. // Original image width, height, shrink in float format
  33. std::vector<float> ori_im_size_f_;
  34. // Evaluation image width and height
  35. std::vector<float> eval_im_size_f_;
  36. // Scale factor for image size to origin image size
  37. std::vector<float> scale_factor_f_;
  38. };
  39. // Abstraction of preprocessing opration class
  40. class PreprocessOp {
  41. public:
  42. virtual void Init(const YAML::Node& item, const std::string& arch) = 0;
  43. virtual void Run(cv::Mat* im, ImageBlob* data) = 0;
  44. };
  45. class InitInfo : public PreprocessOp{
  46. public:
  47. virtual void Init(const YAML::Node& item, const std::string& arch) {}
  48. virtual void Run(cv::Mat* im, ImageBlob* data);
  49. };
  50. class Normalize : public PreprocessOp {
  51. public:
  52. virtual void Init(const YAML::Node& item, const std::string& arch) {
  53. mean_ = item["mean"].as<std::vector<float>>();
  54. scale_ = item["std"].as<std::vector<float>>();
  55. is_channel_first_ = item["is_channel_first"].as<bool>();
  56. is_scale_ = item["is_scale"].as<bool>();
  57. }
  58. virtual void Run(cv::Mat* im, ImageBlob* data);
  59. private:
  60. // CHW or HWC
  61. bool is_channel_first_;
  62. bool is_scale_;
  63. std::vector<float> mean_;
  64. std::vector<float> scale_;
  65. };
  66. class Permute : public PreprocessOp {
  67. public:
  68. virtual void Init(const YAML::Node& item, const std::string& arch) {
  69. to_bgr_ = item["to_bgr"].as<bool>();
  70. is_channel_first_ = item["channel_first"].as<bool>();
  71. }
  72. virtual void Run(cv::Mat* im, ImageBlob* data);
  73. private:
  74. // RGB to BGR
  75. bool to_bgr_;
  76. // CHW or HWC
  77. bool is_channel_first_;
  78. };
  79. class Resize : public PreprocessOp {
  80. public:
  81. virtual void Init(const YAML::Node& item, const std::string& arch) {
  82. arch_ = arch;
  83. interp_ = item["interp"].as<int>();
  84. max_size_ = item["max_size"].as<int>();
  85. if (item["image_shape"].IsDefined()) {
  86. image_shape_ = item["image_shape"].as<std::vector<int>>();
  87. }
  88. target_size_ = item["target_size"].as<int>();
  89. }
  90. // Compute best resize scale for x-dimension, y-dimension
  91. std::pair<float, float> GenerateScale(const cv::Mat& im);
  92. virtual void Run(cv::Mat* im, ImageBlob* data);
  93. private:
  94. std::string arch_;
  95. int interp_;
  96. int max_size_;
  97. int target_size_;
  98. std::vector<int> image_shape_;
  99. };
  100. // Models with FPN need input shape % stride == 0
  101. class PadStride : public PreprocessOp {
  102. public:
  103. virtual void Init(const YAML::Node& item, const std::string& arch) {
  104. stride_ = item["stride"].as<int>();
  105. }
  106. virtual void Run(cv::Mat* im, ImageBlob* data);
  107. private:
  108. int stride_;
  109. };
  110. class Preprocessor {
  111. public:
  112. void Init(const YAML::Node& config_node, const std::string& arch) {
  113. arch_ = arch;
  114. // initialize image info at first
  115. ops_["InitInfo"] = std::make_shared<InitInfo>();
  116. for (const auto& item : config_node) {
  117. auto op_name = item["type"].as<std::string>();
  118. ops_[op_name] = CreateOp(op_name);
  119. ops_[op_name]->Init(item, arch);
  120. }
  121. }
  122. std::shared_ptr<PreprocessOp> CreateOp(const std::string& name) {
  123. if (name == "Resize") {
  124. return std::make_shared<Resize>();
  125. } else if (name == "Permute") {
  126. return std::make_shared<Permute>();
  127. } else if (name == "Normalize") {
  128. return std::make_shared<Normalize>();
  129. } else if (name == "PadStride") {
  130. return std::make_shared<PadStride>();
  131. }
  132. return nullptr;
  133. }
  134. void Run(cv::Mat* im, ImageBlob* data);
  135. public:
  136. static const std::vector<std::string> RUN_ORDER;
  137. private:
  138. std::string arch_;
  139. std::unordered_map<std::string, std::shared_ptr<PreprocessOp>> ops_;
  140. };
  141. } // namespace PaddleDetection