feature_extractor.h 2.4 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667
  1. /*************************************************************************
  2. * Copyright (C) [2019] by Cambricon, Inc. All rights reserved
  3. *
  4. * Licensed under the Apache License, Version 2.0 (the "License");
  5. * you may not use this file except in compliance with the License.
  6. * You may obtain a copy of the License at
  7. *
  8. * http://www.apache.org/licenses/LICENSE-2.0
  9. *
  10. * The above copyright notice and this permission notice shall be included in
  11. * all copies or substantial portions of the Software.
  12. * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
  13. * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  14. * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
  15. * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  16. * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
  17. * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
  18. * THE SOFTWARE.
  19. *************************************************************************/
  20. #ifndef EDK_SAMPLES_FEATURE_EXTRACTOR_H_
  21. #define EDK_SAMPLES_FEATURE_EXTRACTOR_H_
  22. #include <opencv2/core/core.hpp>
  23. #include <memory>
  24. #include <mutex>
  25. #include <string>
  26. #include <utility>
  27. #include <vector>
  28. #include "easyinfer/easy_infer.h"
  29. #include "easyinfer/mlu_memory_op.h"
  30. #include "easyinfer/model_loader.h"
  31. #include "easytrack/easy_track.h"
  32. class FeatureExtractor {
  33. public:
  34. ~FeatureExtractor();
  35. bool Init(const std::string& model_path, const std::string& func_name, int dev_id = 0);
  36. void Destroy();
  37. /*******************************************************
  38. * @brief inference and extract feature of an object
  39. * @param
  40. * frame[in] full image
  41. * obj[in] detected object
  42. * @return return a 128 dimension vector as feature of
  43. * object.
  44. * *****************************************************/
  45. std::vector<float> ExtractFeature(const edk::TrackFrame& frame, const edk::DetectObject& obj);
  46. private:
  47. void Preprocess(const cv::Mat& img);
  48. edk::EasyInfer infer_;
  49. edk::MluMemoryOp mem_op_;
  50. std::shared_ptr<edk::ModelLoader> model_ = nullptr;
  51. std::mutex mlu_proc_mutex_;
  52. int device_id_ = 0;
  53. void** input_cpu_ptr_ = nullptr;
  54. void** output_cpu_ptr_ = nullptr;
  55. void** input_mlu_ptr_ = nullptr;
  56. void** output_mlu_ptr_ = nullptr;
  57. bool extract_feature_mlu_ = false;
  58. }; // class FeatureExtractor
  59. #endif // EDK_SAMPLES_FEATURE_EXTRACTOR_H_