62 lines
1.6 KiB
C++
62 lines
1.6 KiB
C++
|
#include "pybind11/pybind11.h"
|
||
|
|
#include "pybind11/numpy.h"
|
||
|
|
#include "pybind11/stl.h"
|
||
|
|
#include "pybind11/stl_bind.h"
|
||
|
|
|
||
|
|
#include "lanms.h"
|
||
|
|
|
||
|
|
namespace py = pybind11;
|
||
|
|
|
||
|
|
|
||
|
|
namespace lanms_adaptor {
|
||
|
|
|
||
|
|
std::vector<std::vector<float>> polys2floats(const std::vector<lanms::Polygon> &polys) {
|
||
|
|
std::vector<std::vector<float>> ret;
|
||
|
|
for (size_t i = 0; i < polys.size(); i ++) {
|
||
|
|
auto &p = polys[i];
|
||
|
|
auto &poly = p.poly;
|
||
|
|
ret.emplace_back(std::vector<float>{
|
||
|
|
float(poly[0].X), float(poly[0].Y),
|
||
|
|
float(poly[1].X), float(poly[1].Y),
|
||
|
|
float(poly[2].X), float(poly[2].Y),
|
||
|
|
float(poly[3].X), float(poly[3].Y),
|
||
|
|
float(p.score),
|
||
|
|
});
|
||
|
|
}
|
||
|
|
|
||
|
|
return ret;
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
/**
|
||
|
|
*
|
||
|
|
* \param quad_n9 an n-by-9 numpy array, where first 8 numbers denote the
|
||
|
|
* quadrangle, and the last one is the score
|
||
|
|
* \param iou_threshold two quadrangles with iou score above this threshold
|
||
|
|
* will be merged
|
||
|
|
*
|
||
|
|
* \return an n-by-9 numpy array, the merged quadrangles
|
||
|
|
*/
|
||
|
|
std::vector<std::vector<float>> merge_quadrangle_n9(
|
||
|
|
py::array_t<float, py::array::c_style | py::array::forcecast> quad_n9,
|
||
|
|
float iou_threshold) {
|
||
|
|
auto pbuf = quad_n9.request();
|
||
|
|
if (pbuf.ndim != 2 || pbuf.shape[1] != 9)
|
||
|
|
throw std::runtime_error("quadrangles must have a shape of (n, 9)");
|
||
|
|
auto n = pbuf.shape[0];
|
||
|
|
auto ptr = static_cast<float *>(pbuf.ptr);
|
||
|
|
return polys2floats(lanms::merge_quadrangle_n9(ptr, n, iou_threshold));
|
||
|
|
}
|
||
|
|
|
||
|
|
}
|
||
|
|
|
||
|
|
PYBIND11_PLUGIN(adaptor) {
|
||
|
|
py::module m("adaptor", "NMS");
|
||
|
|
|
||
|
|
m.def("merge_quadrangle_n9", &lanms_adaptor::merge_quadrangle_n9,
|
||
|
|
"merge quadrangels");
|
||
|
|
|
||
|
|
return m.ptr();
|
||
|
|
}
|
||
|
|
|