video_postprocess_ssd.cpp 3.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485
  1. /*************************************************************************
  2. * Copyright (C) [2021] 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 <string>
  22. #include <utility>
  23. #include <vector>
  24. #include "cnstream_frame_va.hpp"
  25. #include "cnstream_logging.hpp"
  26. #include "video_postproc.hpp"
  27. /**
  28. * @brief Video postprocessing for ssd neural network
  29. */
  30. class VideoPostprocSsd : public cnstream::VideoPostproc {
  31. public:
  32. /**
  33. * @brief Execute ssd neural network postprocessing
  34. *
  35. * @param output_data: postproc result. The result of postprocessing should be set to it.
  36. * You could set any type of data to this parameter and get it in UserProcess function.
  37. * @param model_output: the raw output data from neural network
  38. * @param model_info: model information, e.g., input/output number, shape and etc.
  39. *
  40. * @return return true if succeed
  41. * @see VideoPostprocSsd::UserProcess
  42. */
  43. bool Execute(infer_server::InferData* output_data, const infer_server::ModelIO& model_output,
  44. const infer_server::ModelInfo& model_info) override;
  45. DECLARE_REFLEX_OBJECT_EX(VideoPostprocSsd, cnstream::VideoPostproc)
  46. }; // class VideoPostprocSsd
  47. IMPLEMENT_REFLEX_OBJECT_EX(VideoPostprocSsd, cnstream::VideoPostproc)
  48. bool VideoPostprocSsd::Execute(infer_server::InferData* output_data,
  49. const infer_server::ModelIO& model_output,
  50. const infer_server::ModelInfo& model_info) {
  51. LOGF_IF(DEMO, model_info.InputNum() != 1);
  52. LOGF_IF(DEMO, model_info.OutputNum() != 1);
  53. LOGF_IF(DEMO, model_output.buffers.size() != 1);
  54. cnstream::CNObjsVec objs;
  55. const float* data = reinterpret_cast<const float*>(model_output.buffers[0].Data());
  56. unsigned box_num = static_cast<unsigned>(data[0]);
  57. data += 64;
  58. for (unsigned bi = 0; bi < box_num; ++bi) {
  59. if (data[1] == 0) continue;
  60. if (threshold_ > 0 && data[2] < threshold_) continue;
  61. std::shared_ptr<cnstream::CNInferObject> object = std::make_shared<cnstream::CNInferObject>();
  62. object->id = std::to_string(static_cast<int>(data[1] - 1));
  63. object->score = data[2];
  64. object->bbox.x = data[3];
  65. object->bbox.y = data[4];
  66. object->bbox.w = data[5] - object->bbox.x;
  67. object->bbox.h = data[6] - object->bbox.y;
  68. objs.push_back(object);
  69. data += 7;
  70. }
  71. cnstream::CNFrameInfoPtr frame = output_data->GetUserData<cnstream::CNFrameInfoPtr>();
  72. cnstream::CNInferObjsPtr objs_holder = frame->collection.Get<cnstream::CNInferObjsPtr>(cnstream::kCNInferObjsTag);
  73. std::lock_guard<std::mutex> objs_mutex(objs_holder->mutex_);
  74. objs_holder->objs_.insert(objs_holder->objs_.end(), objs.begin(), objs.end());
  75. return true;
  76. }