123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110 |
- #include <algorithm>
- #include <cmath>
- #include <memory>
- #include <mutex>
- #include <vector>
- #include "cnstream_frame_va.hpp"
- #include "postproc.hpp"
- #include "cnstream_logging.hpp"
- static auto range_0_1 = [](float num) { return std::max(.0f, std::min(1.0f, num)); };
- class PostprocYolov5 : public cnstream::Postproc {
- public:
- int Execute(const std::vector<float*>& net_outputs, const std::shared_ptr<edk::ModelLoader>& model,
- const std::shared_ptr<cnstream::CNFrameInfo>& package) {
- LOGF_IF(DEMO, model->InputNum() != 1);
- LOGF_IF(DEMO, model->OutputNum() != 1);
- LOGF_IF(DEMO, net_outputs.size() != 1);
- auto input_shape = model->InputShape(0);
- cnstream::CNDataFramePtr frame = package->collection.Get<cnstream::CNDataFramePtr>(cnstream::kCNDataFrameTag);
- const int img_w = frame->width;
- const int img_h = frame->height;
- const int model_input_w = static_cast<int>(input_shape.W());
- const int model_input_h = static_cast<int>(input_shape.H());
- const float* net_output = net_outputs[0];
-
- const float scaling_factors = std::min(1.0 * model_input_w / img_w, 1.0 * model_input_h / img_h);
-
-
-
-
-
-
-
- const int scaled_w = scaling_factors * img_w;
- const int scaled_h = scaling_factors * img_h;
-
- const int box_num = static_cast<int>(net_output[0]);
- int box_step = 7;
- auto range_0_1 = [](float num) { return std::max(.0f, std::min(1.0f, num)); };
- cnstream::CNInferObjsPtr objs_holder =
- package->collection.Get<cnstream::CNInferObjsPtr>(cnstream::kCNInferObjsTag);
- cnstream::CNObjsVec &objs = objs_holder->objs_;
- for (int box_idx = 0; box_idx < box_num; ++box_idx) {
- float left = net_output[64 + box_idx * box_step + 3];
- float right = net_output[64 + box_idx * box_step + 5];
- float top = net_output[64 + box_idx * box_step + 4];
- float bottom = net_output[64 + box_idx * box_step + 6];
-
- left = (left - (model_input_w - scaled_w) / 2) / scaled_w;
- right = (right - (model_input_w - scaled_w) / 2) / scaled_w;
- top = (top - (model_input_h - scaled_h) / 2) / scaled_h;
- bottom = (bottom - (model_input_h - scaled_h) / 2) / scaled_h;
- left = range_0_1(left);
- right = range_0_1(right);
- top = range_0_1(top);
- bottom = range_0_1(bottom);
- auto obj = std::make_shared<cnstream::CNInferObject>();
- obj->id = std::to_string(static_cast<int>(net_output[64 + box_idx * box_step + 1]));
- obj->score = net_output[64 + box_idx * box_step + 2];
- obj->bbox.x = left;
- obj->bbox.y = top;
- obj->bbox.w = std::min(1.0f - obj->bbox.x, right - left);
- obj->bbox.h = std::min(1.0f - obj->bbox.y, bottom - top);
- if (obj->bbox.h <= 0 || obj->bbox.w <= 0 || (obj->score < threshold_ && threshold_ > 0)) continue;
- std::lock_guard<std::mutex> objs_mutex(objs_holder->mutex_);
- objs.push_back(obj);
- }
- return 0;
- }
- private:
- DECLARE_REFLEX_OBJECT_EX(PostprocYolov5, cnstream::Postproc);
- };
- IMPLEMENT_REFLEX_OBJECT_EX(PostprocYolov5, cnstream::Postproc);
|