paddlex.h 2.0 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273
  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 <functional>
  16. #include <iostream>
  17. #include <numeric>
  18. #include "yaml-cpp/yaml.h"
  19. #ifdef _WIN32
  20. #define OS_PATH_SEP "\\"
  21. #else
  22. #define OS_PATH_SEP "/"
  23. #endif
  24. #include "paddle_inference_api.h" // NOLINT
  25. #include "include/paddlex/config_parser.h"
  26. #include "include/paddlex/results.h"
  27. #include "include/paddlex/transforms.h"
  28. namespace PaddleX {
  29. class Model {
  30. public:
  31. void Init(const std::string& model_dir,
  32. bool use_gpu = false,
  33. bool use_trt = false,
  34. int gpu_id = 0) {
  35. create_predictor(model_dir, use_gpu, use_trt, gpu_id);
  36. }
  37. void create_predictor(const std::string& model_dir,
  38. bool use_gpu = false,
  39. bool use_trt = false,
  40. int gpu_id = 0);
  41. bool load_config(const std::string& model_dir);
  42. bool preprocess(const cv::Mat& input_im, ImageBlob* blob);
  43. bool predict(const cv::Mat& im, ClsResult* result);
  44. bool predict(const cv::Mat& im, DetResult* result);
  45. bool predict(const cv::Mat& im, SegResult* result);
  46. bool postprocess(SegResult* result);
  47. bool postprocess(DetResult* result);
  48. std::string type;
  49. std::string name;
  50. std::map<int, std::string> labels;
  51. Transforms transforms_;
  52. ImageBlob inputs_;
  53. std::vector<float> outputs_;
  54. std::unique_ptr<paddle::PaddlePredictor> predictor_;
  55. };
  56. } // namespce of PaddleX