preprocess_op.h 4.6 KB

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