transforms.h 3.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111
  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 <memory>
  17. #include <string>
  18. #include <unordered_map>
  19. #include <utility>
  20. #include <vector>
  21. #include <opencv2/core/core.hpp>
  22. #include <opencv2/highgui/highgui.hpp>
  23. #include <opencv2/imgproc/imgproc.hpp>
  24. #include <inference_engine.hpp>
  25. using namespace InferenceEngine;
  26. namespace PaddleX {
  27. // Abstraction of preprocessing opration class
  28. class Transform {
  29. public:
  30. virtual void Init(const YAML::Node& item) = 0;
  31. virtual bool Run(cv::Mat* im, ImageBlob* data) = 0;
  32. };
  33. class Normalize : public Transform {
  34. public:
  35. virtual void Init(const YAML::Node& item) {
  36. mean_ = item["mean"].as<std::vector<float>>();
  37. std_ = item["std"].as<std::vector<float>>();
  38. }
  39. virtual bool Run(cv::Mat* im, ImageBlob* data);
  40. private:
  41. std::vector<float> mean_;
  42. std::vector<float> std_;
  43. };
  44. class Resize : public Transform {
  45. public:
  46. virtual void Init(const YAML::Node& item) {
  47. if (item["target_size"].IsScalar()) {
  48. height_ = item["target_size"].as<int>();
  49. width_ = item["target_size"].as<int>();
  50. interp_ = item["interp"].as<std::string>();
  51. } else if (item["target_size"].IsSequence()) {
  52. std::vector<int> target_size = item["target_size"].as<std::vector<int>>();
  53. width_ = target_size[0];
  54. height_ = target_size[1];
  55. }
  56. if (height_ <= 0 || width_ <= 0) {
  57. std::cerr << "[Resize] target_size should greater than 0" << std::endl;
  58. exit(-1);
  59. }
  60. }
  61. virtual bool Run(cv::Mat* im, ImageBlob* data);
  62. private:
  63. int height_;
  64. int width_;
  65. std::string interp_;
  66. };
  67. class CenterCrop : public Transform {
  68. public:
  69. virtual void Init(const YAML::Node& item) {
  70. if (item["crop_size"].IsScalar()) {
  71. height_ = item["crop_size"].as<int>();
  72. width_ = item["crop_size"].as<int>();
  73. } else if (item["crop_size"].IsSequence()) {
  74. std::vector<int> crop_size = item["crop_size"].as<std::vector<int>>();
  75. width_ = crop_size[0];
  76. height_ = crop_size[1];
  77. }
  78. }
  79. virtual bool Run(cv::Mat* im, ImageBlob* data);
  80. private:
  81. int height_;
  82. int width_;
  83. };
  84. class Transforms {
  85. public:
  86. void Init(const YAML::Node& node, bool to_rgb = true);
  87. std::shared_ptr<Transform> CreateTransform(const std::string& name);
  88. bool Run(cv::Mat* im, Blob::ptr data);
  89. private:
  90. std::vector<std::shared_ptr<Transform>> transforms_;
  91. bool to_rgb_ = true;
  92. };
  93. } // namespace PaddleX