cnstream/python/test/cpp_postproc_test.cpp

#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

#include <memory>
#include <functional>
#include <string>
#include <unordered_map>
#include <vector>

#include "pypostproc.h"

namespace py = pybind11;

bool TestPyPostproc(const std::unordered_map<std::string, std::string> &params) {
  auto model = std::make_shared<edk::ModelLoader>("data/test_model.cambricon", "subnet0");
  std::vector<float*> outputs;
  for (uint32_t i = 0; i < model->OutputNum(); ++i) {
    outputs.push_back(new float[model->OutputShape(i).DataCount()]);
  }
  auto free_outputs = [&outputs] () {
    for (auto ptr : outputs) delete[] ptr;
  };
  cnstream::PyPostproc pypostproc;
  if (!pypostproc.Init(params)) {
    free_outputs();
    return false;
  }
  bool ret = 0 == pypostproc.Execute(outputs, model, nullptr);
  free_outputs();
  return ret;
}

bool TestPyObjPostproc(const std::unordered_map<std::string, std::string> &params) {
  auto model = std::make_shared<edk::ModelLoader>("data/test_model.cambricon", "subnet0");
  std::vector<float*> outputs;
  for (uint32_t i = 0; i < model->OutputNum(); ++i) {
    outputs.push_back(new float[model->OutputShape(i).DataCount()]);
  }
  auto free_outputs = [&outputs] () {
    for (auto ptr : outputs) delete[] ptr;
  };
  cnstream::PyObjPostproc pyobjpostproc;
  if (!pyobjpostproc.Init(params)) {
    free_outputs();
    return false;
  }
  bool ret = 0 == pyobjpostproc.Execute(outputs, model, nullptr, nullptr);
  free_outputs();
  return ret;
}

void PostprocTestWrapper(py::module &m) {  // NOLINT
  m.def("cpptest_pypostproc", &TestPyPostproc);
  m.def("cpptest_pyobjpostproc", &TestPyObjPostproc);
}