preprocess_yolov3.cpp 3.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778
  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. #include <algorithm>
  21. #include <memory>
  22. #include <vector>
  23. #include "opencv2/highgui/highgui.hpp"
  24. #include "opencv2/imgproc/imgproc.hpp"
  25. #if (CV_MAJOR_VERSION >= 3)
  26. #include "opencv2/imgcodecs/imgcodecs.hpp"
  27. #endif
  28. #include "cnstream_frame_va.hpp"
  29. #include "preproc.hpp"
  30. #include "cnstream_logging.hpp"
  31. class PreprocYolov3 : public cnstream::Preproc {
  32. public:
  33. int Execute(const std::vector<float*>& net_inputs, const std::shared_ptr<edk::ModelLoader>& model,
  34. const std::shared_ptr<cnstream::CNFrameInfo>& package) {
  35. // check params
  36. auto input_shape = model->InputShape(0);
  37. if (net_inputs.size() != 1 || input_shape.C() != 3) {
  38. LOGE(DEMO) << "[PreprocCpu] model input shape not supported";
  39. return -1;
  40. }
  41. cnstream::CNDataFramePtr frame = package->collection.Get<cnstream::CNDataFramePtr>(cnstream::kCNDataFrameTag);
  42. int width = frame->width;
  43. int height = frame->height;
  44. int dst_w = input_shape.W();
  45. int dst_h = input_shape.H();
  46. cv::Mat img = frame->ImageBGR();
  47. // resize
  48. if (height != dst_h || width != dst_w) {
  49. cv::Mat dst(dst_h, dst_w, CV_8UC3, cv::Scalar(128, 128, 128));
  50. const float scaling_factors = std::min(1.0 * dst_w / width, 1.0 * dst_h / height);
  51. cv::Mat resized(height * scaling_factors, width * scaling_factors, CV_8UC3);
  52. cv::resize(img, resized, cv::Size(resized.cols, resized.rows));
  53. cv::Rect roi;
  54. roi.x = (dst.cols - resized.cols) / 2;
  55. roi.y = (dst.rows - resized.rows) / 2;
  56. roi.width = resized.cols;
  57. roi.height = resized.rows;
  58. resized.copyTo(dst(roi));
  59. img = dst;
  60. }
  61. // since model input data type is float, convert image to float
  62. cv::Mat dst(dst_h, dst_w, CV_32FC3, net_inputs[0]);
  63. img.convertTo(dst, CV_32F);
  64. return 0;
  65. }
  66. private:
  67. DECLARE_REFLEX_OBJECT_EX(PreprocYolov3, cnstream::Preproc);
  68. }; // class PreprocYolov3
  69. IMPLEMENT_REFLEX_OBJECT_EX(PreprocYolov3, cnstream::Preproc);