preprocess_standard.cpp 6.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185
  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 <memory>
  21. #include <utility>
  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 "easyinfer/model_loader.h"
  30. #include "easyinfer/shape.h"
  31. #include "preproc.hpp"
  32. #include "cnstream_logging.hpp"
  33. /**
  34. * @brief standard pre process
  35. */
  36. class PreprocCpu : public cnstream::Preproc {
  37. public:
  38. /**
  39. * @brief Execute preproc on origin data
  40. *
  41. * @param net_inputs: neural network inputs
  42. * @param model: model information(you can get input shape and output shape from model)
  43. * @param package: smart pointer of struct to store origin data
  44. *
  45. * @return return 0 if succeed
  46. *
  47. * @attention net_inputs is a pointer to pre-allocated cpu memory
  48. */
  49. int Execute(const std::vector<float*>& net_inputs, const std::shared_ptr<edk::ModelLoader>& model,
  50. const cnstream::CNFrameInfoPtr& package) override;
  51. DECLARE_REFLEX_OBJECT_EX(PreprocCpu, cnstream::Preproc);
  52. }; // class PreprocCpu
  53. IMPLEMENT_REFLEX_OBJECT_EX(PreprocCpu, cnstream::Preproc)
  54. int PreprocCpu::Execute(const std::vector<float*>& net_inputs, const std::shared_ptr<edk::ModelLoader>& model,
  55. const cnstream::CNFrameInfoPtr& package) {
  56. // check params
  57. edk::ShapeEx input_shape;
  58. try {
  59. input_shape = model->InputShape(0);
  60. if (net_inputs.size() != 1 || (input_shape.C() != 3 && input_shape.C() != 4)) {
  61. LOGE(DEMO) << "[PreprocCpu] model input shape not supported, net_input.size = " << net_inputs.size()
  62. << ", input_shape.c = " << input_shape.C();
  63. return -1;
  64. }
  65. } catch (const edk::Exception& e) {
  66. LOGE(DEMO) << e.what();
  67. return -1;
  68. }
  69. LOGI(DEMO) << "[PreprocCpu] do preproc...";
  70. cnstream::CNDataFramePtr frame = package->collection.Get<cnstream::CNDataFramePtr>(cnstream::kCNDataFrameTag);
  71. int width = frame->width;
  72. int height = frame->height;
  73. int dst_w = input_shape.W();
  74. int dst_h = input_shape.H();
  75. uint8_t* img_data = new (std::nothrow) uint8_t[frame->GetBytes()];
  76. if (!img_data) {
  77. LOGE(DEMO) << "Failed to alloc memory, size: " << frame->GetBytes();
  78. return -1;
  79. }
  80. uint8_t* t = img_data;
  81. for (int i = 0; i < frame->GetPlanes(); ++i) {
  82. memcpy(t, frame->data[i]->GetCpuData(), frame->GetPlaneBytes(i));
  83. t += frame->GetPlaneBytes(i);
  84. }
  85. // convert color space
  86. cv::Mat img;
  87. switch (frame->fmt) {
  88. case cnstream::CNDataFormat::CN_PIXEL_FORMAT_BGR24:
  89. img = cv::Mat(height, width, CV_8UC3, img_data);
  90. break;
  91. case cnstream::CNDataFormat::CN_PIXEL_FORMAT_RGB24:
  92. img = cv::Mat(height, width, CV_8UC3, img_data);
  93. cv::cvtColor(img, img, cv::COLOR_RGB2BGR);
  94. break;
  95. case cnstream::CNDataFormat::CN_PIXEL_FORMAT_YUV420_NV12: {
  96. img = cv::Mat(height * 3 / 2, width, CV_8UC1, img_data);
  97. cv::Mat bgr(height, width, CV_8UC3);
  98. cv::cvtColor(img, bgr, cv::COLOR_YUV2BGR_NV12);
  99. img = bgr;
  100. } break;
  101. case cnstream::CNDataFormat::CN_PIXEL_FORMAT_YUV420_NV21: {
  102. img = cv::Mat(height * 3 / 2, width, CV_8UC1, img_data);
  103. cv::Mat bgr(height, width, CV_8UC3);
  104. cv::cvtColor(img, bgr, cv::COLOR_YUV2BGR_NV21);
  105. img = bgr;
  106. } break;
  107. default:
  108. LOGW(DEMO) << "[PreprocCpu] Unsupport pixel format.";
  109. delete[] img_data;
  110. return -1;
  111. }
  112. // resize if needed
  113. if (height != dst_h || width != dst_w) {
  114. cv::Mat dst(dst_h, dst_w, CV_8UC3);
  115. cv::resize(img, dst, cv::Size(dst_w, dst_h));
  116. img.release();
  117. img = dst;
  118. }
  119. // since model input data type is float, convert image to float
  120. cv::Mat dst(dst_h, dst_w, CV_32FC3, net_inputs[0]);
  121. img.convertTo(dst, CV_32F);
  122. delete[] img_data;
  123. return 0;
  124. }
  125. /**
  126. * @brief standard object pre process
  127. */
  128. class ObjPreprocCpu : public cnstream::ObjPreproc {
  129. public:
  130. int Execute(const std::vector<float*>& net_inputs, const std::shared_ptr<edk::ModelLoader>& model,
  131. const cnstream::CNFrameInfoPtr& finfo, const std::shared_ptr<cnstream::CNInferObject>& pobj) override;
  132. DECLARE_REFLEX_OBJECT_EX(ObjPreprocCpu, cnstream::ObjPreproc);
  133. }; // class ObjPreprocCpu
  134. IMPLEMENT_REFLEX_OBJECT_EX(ObjPreprocCpu, cnstream::ObjPreproc)
  135. int ObjPreprocCpu::Execute(const std::vector<float*>& net_inputs, const std::shared_ptr<edk::ModelLoader>& model,
  136. const cnstream::CNFrameInfoPtr& finfo,
  137. const std::shared_ptr<cnstream::CNInferObject>& pobj) {
  138. cnstream::CNDataFramePtr frame = finfo->collection.Get<cnstream::CNDataFramePtr>(cnstream::kCNDataFrameTag);
  139. // origin frame
  140. cv::Mat frame_bgr = frame->ImageBGR();
  141. // crop objct from frame
  142. int w = frame->width;
  143. int h = frame->height;
  144. cv::Rect obj_roi(pobj->bbox.x * w, pobj->bbox.y * h, pobj->bbox.w * w, pobj->bbox.h * h);
  145. cv::Mat obj_bgr = frame_bgr(obj_roi);
  146. // resize
  147. int input_w = model->InputShape(0).W();
  148. int input_h = model->InputShape(0).H();
  149. cv::Mat obj_bgr_resized;
  150. cv::resize(obj_bgr, obj_bgr_resized, cv::Size(input_w, input_h));
  151. // bgr2bgra
  152. cv::Mat obj_bgra;
  153. cv::Mat a(input_h, input_w, CV_8UC1, cv::Scalar(0.0));
  154. std::vector<cv::Mat> vec_mat = {obj_bgr_resized, a};
  155. cv::merge(std::move(vec_mat), obj_bgra);
  156. // convert to float32, required by inferencer module
  157. cv::Mat obj_bgra_float32(input_h, input_w, CV_32FC4, net_inputs[0]);
  158. obj_bgra.convertTo(obj_bgra_float32, CV_32FC4);
  159. return 0;
  160. }